Update README.md

README.md

@@ -55,19 +55,48 @@ README.md
 ```python
 from sentence_transformers import SentenceTransformer
 import xgboost as xgb
-
-embedder = SentenceTransformer("sentence_transformer_model")
-
+import numpy as np
+from huggingface_hub import hf_hub_download
+
+# ----------------------------------------------------
+# 1. Load the embedding model (must be downloaded locally)
+# ----------------------------------------------------
+# Users must install: pip install sentence-transformers
+embedder = SentenceTransformer("all-MiniLM-L6-v2")
+
+# ----------------------------------------------------
+# 2. Download your XGBoost model from HuggingFace Hub
+# ----------------------------------------------------
+model_path = hf_hub_download(
+    repo_id="mjpsm/checkin_or_not_model",
+    filename="checkin_or_not_classifier.json"
+)
+
+# Load the model
 booster = xgb.Booster()
-booster.load_model("checkin_or_not_classifier.json")
+booster.load_model(model_path)
 
-def predict(text):
+# ----------------------------------------------------
+# 3. Prediction function
+# ----------------------------------------------------
+def predict(text: str):
     emb = embedder.encode([text])
-    pred = booster.predict(xgb.DMatrix(emb))[0]
-    return ("CHECKIN" if pred >= 0.5 else "NOT_CHECKIN", float(pred))
+    dmatrix = xgb.DMatrix(emb)
+
+    score = float(booster.predict(dmatrix)[0])
+    label = "CHECKIN" if score >= 0.5 else "NOT_CHECKIN"
+
+    return {"label": label, "score": score}
+
+
+# ----------------------------------------------------
+# 4. Example usage
+# ----------------------------------------------------
+example = "Today I worked on improving the automation workflow."
+result = predict(example)
+
+print(result)
 
-label, score = predict("Today I worked on automation tasks.")
-print(label, score)
 ```
 
 ## 📘 Intended Use