Update app.py

app.py

@@ -2,76 +2,64 @@ import json
 import numpy as np
 from fastapi import FastAPI
 from pydantic import BaseModel
+from sentence_transformers import SentenceTransformer
 from huggingface_hub import hf_hub_download
-from transformers import AutoTokenizer, AutoModel
-import torch
 
-# -----------------------------
+# ---------------------------------------------------------
 # CONFIG
-# -----------------------------
-EMBEDDER_REPO = "ClergeF/MVT-embedder"
-MODEL_REPO = "ClergeF/impact-model"
-MODEL_FILE = "impact.json"
+# ---------------------------------------------------------
+HF_USER = "ClergeF"
+IMPACT_REPO = "impact-model"
+IMPACT_FILE = "impact.json"
 
-# -----------------------------
-# SAFE EMBEDDER LOADER
-# -----------------------------
-def load_safe_embedder(repo_id: str):
-    print(f"Loading embedder from {repo_id} (safe mode)...")
-    tokenizer = AutoTokenizer.from_pretrained(repo_id)
-    model = AutoModel.from_pretrained(repo_id)
+# ---------------------------------------------------------
+# LOAD EMBEDDER
+# ---------------------------------------------------------
+print("Loading embedder: all-MiniLM-L6-v2 ...")
+embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
-    def embed_fn(texts):
-        tokens = tokenizer(
-            texts,
-            padding=True,
-            truncation=True,
-            max_length=256,
-            return_tensors="pt"
-        )
-        with torch.no_grad():
-            outputs = model(**tokens).last_hidden_state
-            embeddings = outputs.mean(dim=1).numpy()
-        return embeddings
+# ---------------------------------------------------------
+# LOAD IMPACT MODEL
+# ---------------------------------------------------------
+def load_json_model(repo_id, filename):
+    """Download model.json from HuggingFace Hub"""
+    path = hf_hub_download(repo_id=repo_id, filename=filename)
+    with open(path, "r") as f:
+        return json.load(f)
 
-    return embed_fn
+print("Loading impact model...")
+impact_model = load_json_model(f"{HF_USER}/{IMPACT_REPO}", IMPACT_FILE)
 
-# Load embedder
-embed_fn = load_safe_embedder(EMBEDDER_REPO)
+# ---------------------------------------------------------
+# FASTAPI
+# ---------------------------------------------------------
+app = FastAPI(title="Impact Level Model API")
 
-# -----------------------------
-# LOAD THE IMPACT MODEL
-# -----------------------------
-print("Loading impact model...")
-path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
-with open(path, "r") as f:
-    impact_model = json.load(f)
+class InputText(BaseModel):
+    text: str
 
-# -----------------------------
-# HELPERS
-# -----------------------------
+# ---------------------------------------------------------
+# PREDICT HELPERS
+# ---------------------------------------------------------
 def embed(text: str):
-    return embed_fn([text])[0]
+    return embedder.encode([text])[0]
 
 def linear_predict(model_json, vec):
     coef = np.array(model_json["coef"])
-    intercept = np.array(model_json["intercept"])
+    intercept = float(model_json["intercept"])
     return float(np.dot(coef, vec) + intercept)
 
-# -----------------------------
-# FASTAPI
-# -----------------------------
-app = FastAPI(title="Impact Model API")
-
-class Input(BaseModel):
-    text: str
-
+# ---------------------------------------------------------
+# ROUTE
+# ---------------------------------------------------------
 @app.post("/predict")
-def predict(payload: Input):
-    vec = embed(payload.text)
-    impact = linear_predict(impact_model, vec)
+def predict(payload: InputText):
+    text = payload.text
+    vec = embed(text)
+
+    impact_level = linear_predict(impact_model, vec)
 
     return {
-        "input": payload.text,
-        "impact_score": impact
+        "input": text,
+        "impact_score": impact_level
     }