Update app.py

app.py

@@ -22,7 +22,6 @@ base = AutoModelForCausalLM.from_pretrained(
 model = PeftModel.from_pretrained(base, ADAPTER_REPO).to(device)
 model.eval()
 
-# Optional merge (can speed up). If it fails, just continue.
 try:
     model = model.merge_and_unload()
     model.to(device)
@@ -30,11 +29,18 @@ try:
 except Exception:
     pass
 
+LABELS = ["negative", "neutral", "positive"]
+
+label_token_ids = {
+    lab: tokenizer(" " + lab, add_special_tokens=False)["input_ids"]
+    for lab in LABELS
+}
+
 finance_words = [
     "stock","shares","profit","profits","loss","losses","revenue","earnings","dividend","market",
     "bank","loan","interest","inflation","bond","equity","merger","acquisition",
     "ipo","valuation","cash","cashflow","forecast","guidance","quarter","q1","q2","q3","q4",
-    "ceo","cfo","board","layoffs","bankruptcy","debt","default","margin"
+    "ceo","cfo","board","layoffs","bankruptcy","debt","default","margin","miss","downgrade"
 ]
 
 def looks_finance(text: str) -> bool:
@@ -45,33 +51,29 @@ def is_greeting(text: str) -> bool:
     t = (text or "").lower().strip()
     return t in ["hi", "hello", "hey", "good morning", "good afternoon", "good evening"]
 
-def extract_label(gen_text: str) -> str:
-    """
-    Extract the first occurrence of one of the labels from generated text only.
-    """
-    t = (gen_text or "").lower()
-    m = re.search(r"\b(negative|neutral|positive)\b", t)
-    return m.group(1) if m else "neutral"
-
 @torch.inference_mode()
-def generate_answer_only(prompt: str, max_new_tokens: int = 4) -> str:
+def score_label_with_cache(prompt_ids, lab_ids) -> float:
     """
-    Generate ONLY the new tokens after the prompt.
-    This avoids the 'prompt contains the labels' bug.
+    Score P(label | prompt) using cached past_key_values.
+    Returns average log-prob per label token (length-normalized).
     """
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    input_len = inputs["input_ids"].shape[1]
-
-    out = model.generate(
-        **inputs,
-        max_new_tokens=max_new_tokens,
-        do_sample=False,
-        temperature=0.0,
-        pad_token_id=tokenizer.eos_token_id
-    )
+    # Run prompt once to get cache
+    prompt = torch.tensor([prompt_ids], device=device)
+    out = model(input_ids=prompt, use_cache=True)
+    past = out.past_key_values
+
+    logp_sum = 0.0
+    prev_token = prompt[:, -1:]  
+
+    for tok_id in lab_ids:
+        step = model(input_ids=prev_token, past_key_values=past, use_cache=True)
+        logits = step.logits[:, -1, :] 
+        logp_sum += torch.log_softmax(logits, dim=-1)[0, tok_id].item()
 
-    gen_tokens = out[0][input_len:]
-    return tokenizer.decode(gen_tokens, skip_special_tokens=True).strip()
+        past = step.past_key_values
+        prev_token = torch.tensor([[tok_id]], device=device)
+
+    return logp_sum / max(len(lab_ids), 1)
 
 @torch.inference_mode()
 def predict_label(msg: str) -> str:
@@ -80,11 +82,10 @@ def predict_label(msg: str) -> str:
         f"Text: {msg.strip()}\n"
         "Answer:"
     )
+    prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
 
-    gen = generate_answer_only(prompt, max_new_tokens=4)
-    label = extract_label(gen)
-
-    return label
+    scores = {lab: score_label_with_cache(prompt_ids, label_token_ids[lab]) for lab in LABELS}
+    return max(scores, key=scores.get)
 
 def chat(msg, history):
     msg = (msg or "").strip()