update readme file

README.md

@@ -170,34 +170,62 @@ If included in the repo, the plot `f1_vs_freq.png` shows F1 vs. training label f
 ### Quick start (load from Hub, inference)
 
 ```python
+# pip install sentence-transformers huggingface_hub tensorflow==2.*  # if needed
+import json
 import numpy as np
+from typing import List, Tuple
 from huggingface_hub import hf_hub_download
 from tensorflow import keras
-import json
+from sentence_transformers import SentenceTransformer
 
 REPO_ID = "Bocklitz-Lab/lit2vec-subfield-classifier-model"
+EMBED_MODEL = "intfloat/e5-large-v2"   # must match what you used to train!
+TEXT_PREFIX = {"abstract": "abstract: ", "summary": "summary: "}  # keep consistent with your pipeline
+THRESHOLD = 0.5  # decision threshold for multilabel
 
-# Download files
+# ----- Load model + label mapping -----
 model_path = hf_hub_download(REPO_ID, filename="mlp_model.h5")
 label_map_path = hf_hub_download(REPO_ID, filename="label_mapping.json")
 
-# Load label mapping
-with open(label_map_path) as f:
-    m = json.load(f)
-index_to_label = {int(k): v for k, v in m["index_to_label"].items()}
-
-# Load model (no need for custom loss at inference)
-model = keras.models.load_model(model_path, compile=False)
-
-# x must be a single embedding vector with the correct dimension
-x = np.load("example_embedding.npy").astype("float32")  # shape (D,)
-x = x[None, :]  # shape (1, D)
+with open(label_map_path, "r", encoding="utf-8") as f:
+    mapping = json.load(f)
+index_to_label = {int(k): v for k, v in mapping["index_to_label"].items()}
+
+model = keras.models.load_model(model_path, compile=False)  # inference only
+encoder = SentenceTransformer(EMBED_MODEL)
+
+def encode_text(text: str, text_type: str = "summary") -> np.ndarray:
+    """
+    Encode text into a normalized embedding compatible with the classifier.
+    text_type: "summary" or "abstract" (affects prefix)
+    """
+    prefix = TEXT_PREFIX.get(text_type, "")
+    emb = encoder.encode([prefix + text], normalize_embeddings=True)  # shape: (1, D)
+    return emb.astype("float32")
+
+def predict_labels_from_text(text: str, text_type: str = "summary", threshold: float = THRESHOLD
+                            ) -> Tuple[List[int], List[str], np.ndarray]:
+    """
+    Returns (predicted_ids, predicted_labels, probabilities)
+    """
+    x = encode_text(text, text_type=text_type)       # (1, D)
+    probs = model.predict(x, verbose=0)[0]           # (18,)
+    pred_ids = [i for i, p in enumerate(probs) if p > threshold]
+    pred_labels = [index_to_label[i] for i in pred_ids]
+    return pred_ids, pred_labels, probs
+
+# ----- Example -----
+if __name__ == "__main__":
+    sample_text = (
+        "We investigate Pt-based nanoparticle catalysts for hydrogen evolution, "
+        "demonstrating improved activity via ligand tuning and support engineering."
+    )
+    ids, labels, probs = predict_labels_from_text(sample_text, text_type="abstract", threshold=0.5)
+    print("Predicted IDs:", ids)
+    print("Predicted Labels:", labels)
+    print("Top scores:", sorted(((index_to_label[i], float(p)) for i, p in enumerate(probs)),
+                               key=lambda x: x[1], reverse=True)[:5])
 
-probs = model.predict(x, verbose=0)[0]
-threshold = 0.5
-pred_ids = [i for i, p in enumerate(probs) if p > threshold]
-pred_labels = [index_to_label[i] for i in pred_ids]
-print(list(zip(pred_labels, probs[pred_ids])))
 ```
 
 ### Batch inference