Update app.py

app.py

@@ -1,5 +1,5 @@
 import os
-
+import re
 # Đặt biến cache phòng khi runtime override (khớp Dockerfile)
 os.environ.setdefault("HF_HOME", "/data/hf")
 os.environ.setdefault("HF_HUB_CACHE", "/data/hf/hub")
@@ -43,7 +43,29 @@ def softmax_logs(d):
     ex = {k: math.exp(v - m) for k, v in d.items()}
     Z = sum(ex.values())
     return {k: ex[k]/Z for k in ex}
-
+def generate_text(prompt: str, max_new_tokens: int = 64):
+    # Nếu model cần prefix tác vụ, thêm tại đây, ví dụ:
+    # prompt = f"hate-spans-detection: {prompt}"
+    enc = tok(prompt, return_tensors="pt", truncation=True, max_length=512).to(device)
+    with torch.no_grad():
+        out = model.generate(
+            **enc,
+            max_new_tokens=max_new_tokens,
+            num_beams=4,
+            do_sample=False,
+            early_stopping=True
+        )
+    return tok.decode(out[0], skip_special_tokens=True)
+def extract_hate_spans(output_text: str):
+    # Hỗ trợ cả 2 kiểu: [hate]... [hate]  hoặc  [hate]...[/hate]
+    spans = []
+    # Kiểu 1: [hate]... [hate]
+    spans += re.findall(r"\[hate\](.*?)\[hate\]", output_text, flags=re.IGNORECASE|re.DOTALL)
+    # Kiểu 2: [hate]...[/hate]
+    spans += re.findall(r"\[hate\](.*?)\[/hate\]", output_text, flags=re.IGNORECASE|re.DOTALL)
+    # Làm sạch
+    spans = [s.strip() for s in spans if s.strip()]
+    return spans
 @app.get("/health")
 def health():
     return {"status": "ok", "device": str(device)}
@@ -61,7 +83,15 @@ def predict(item: Item):
     hsd_probs = softmax_logs(score_labels(enc, hsd_labels))
     hsd_label = max(hsd_probs, key=hsd_probs.get)
 
+    span_prompt = text
+    gen = generate_text(span_prompt, max_new_tokens=64)
+    spans = extract_hate_spans(gen)
+
     return {
         "toxic-speech-detection": {"label": tox_label, "probs": tox_probs},
         "hate-speech-detection": {"label": hsd_label, "probs": hsd_probs},
+        "hate-spans-detection": {
+            "spans": spans if spans else [],
+            "raw": gen  # giữ nguyên đầu ra để bạn debug định dạng
+        }
     }