Add compute_score

app.py

@@ -12,7 +12,7 @@ import spaces
 import gradio as gr
 import torch
 import torch.nn.functional as F
-from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer, TextIteratorStreamer, AutoModel
+from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer, TextIteratorStreamer, AutoModel, AutoModelForSequenceClassification
 from kernels import get_kernel
 from typing import Any, Optional, Dict
 
@@ -43,7 +43,7 @@ def load_model(model_id: str):
     print(f"🔹 Loading model: {model_id}")
     tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-    model = AutoModel.from_pretrained(model_id, torch_dtype=torch.bfloat16)
+    model = AutoModelForSequenceClassification.from_pretrained(model_id, torch_dtype=torch.bfloat16)
     print(f"Detected embedding model: {model_id}")
 
     model.to(DEVICE).eval()
@@ -52,11 +52,22 @@ def load_model(model_id: str):
     return model, tokenizer
 
 
-# --- Helper for embeddings ---
-def average_pool(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
-    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
-    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+def compute_score(text: str, model: torch.nn.Module, tokenizer: AutoTokenizer) -> dict:
+        inputs = tokenizer(
+            text,
+            return_tensors="pt",
+            padding="longest",
+            truncation=True,
+        ).to(DEVICE)
 
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logits = outputs.logits.squeeze(-1).float().cpu().numpy()
+
+        res = {}
+        res["score"] = logits.tolist()
+        res["int_score"] = [int(round(max(0, min(score, 5)))) for score in logits]
+        return res
 
 # --- Main scoring logic ---
 @spaces.GPU
@@ -66,21 +77,9 @@ def bot(user_message: str, history: list[dict[str, Any]], model_choice: str):
 
     model, tokenizer = load_model(model_choice)  # returns embedding model
     history = history + [{"role": "user", "content": user_message}]
+    score = compute_score(user_message, model, tokenizer)["score"]
 
-    batch = tokenizer([user_message], padding=True, truncation=True, return_tensors="pt").to(DEVICE)
-    with torch.no_grad():
-        outputs = model(**batch)
-        # outputs.last_hidden_state.shape = [batch_size, seq_len, hidden_dim]
-        # average pool over tokens
-        embedding = average_pool(outputs.last_hidden_state, batch["attention_mask"])
-        score = model.score_head(embedding).squeeze().item()
-        # embedding = F.normalize(embedding, p=2, dim=1)  # optional
-        #
-        # # Compute scalar score from embedding (example: mean of embedding dims)
-        # score = embedding.mean().item()
-
-    response = f"🔹 {model_choice} → score: {score:.4f}"
-    history.append({"role": "assistant", "content": response})
+    history.append({"role": "assistant", "content": f"{model_choice}: {score}"})
     return "", history
 
 # --- UI ---