Spaces:
Sleeping
Sleeping
File size: 2,484 Bytes
cd7b478 76f0f5c cd7b478 e857b33 cd7b478 a12930f cd7b478 a12930f cd7b478 a12930f cd7b478 a12930f cd7b478 a12930f cd7b478 a12930f cd7b478 a12930f cd7b478 a12930f cd7b478 a12930f cd7b478 a12930f 9c04c5d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import re
from fastapi import FastAPI
from pydantic import BaseModel
# ================== CONFIG ==================
MODEL_ID = "AkCh12/claim-extractor-deberta-v3"
MAX_LENGTH = 128
DEFAULT_THRESHOLD = 0.6
device = "cpu"
# ================== LOAD MODEL ==================
print("[INFO] Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
print("[INFO] Loading model...")
model = AutoModelForSequenceClassification.from_pretrained(
MODEL_ID,
num_labels=2,
torch_dtype=torch.float32
)
model.to(device)
model.eval()
torch.set_num_threads(1)
print("[INFO] Warming up model...")
with torch.no_grad():
dummy = tokenizer("Warmup text.", return_tensors="pt")
_ = model(**dummy)
print("[INFO] Model ready.")
# ================== HELPERS ==================
def split_sentences(text: str):
return re.split(r'(?<=[.!?])\s+', text)
def score_sentences(sentences):
inputs = tokenizer(
sentences,
padding=True,
truncation=True,
max_length=MAX_LENGTH,
return_tensors="pt"
)
with torch.no_grad():
logits = model(**inputs).logits
probs = F.softmax(logits, dim=-1)
return probs[:, 1].tolist()
def extract_claims(text: str, threshold: float):
if not text or len(text.strip()) < 20:
return []
sentences = [s.strip() for s in split_sentences(text) if len(s.strip()) > 10]
if not sentences:
return []
scores = score_sentences(sentences)
return [
{
"sentence": s,
"claim_probability": round(sc, 4),
"is_claim": sc >= threshold
}
for s, sc in zip(sentences, scores)
]
# ================== FASTAPI ==================
app = FastAPI(title="Claim Extraction API")
class AnalyzeRequest(BaseModel):
text: str
threshold: float | None = None
@app.post("/extract")
def extract_api(payload: AnalyzeRequest):
threshold = payload.threshold or DEFAULT_THRESHOLD
claims = extract_claims(payload.text, threshold)
return {
"threshold": threshold,
"num_sentences": len(claims),
"claims": claims
}
@app.get("/")
def home():
return {
"message": "Claim Extraction API running",
"model": MODEL_ID
}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860) |