File size: 1,209 Bytes
ad06508 55298f1 ad06508 55298f1 ad06508 | 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 | """
Sinhala-English Codemix PII Detector — HuggingFace Space API
Model: Ramosh99/mixcode_pii_detector (XLM-RoBERTa token-classification)
Also used as fallback for pure Sinhala inputs.
"""
import os
from transformers import pipeline
from fastapi import FastAPI
from pydantic import BaseModel
app = FastAPI(title="Codemix PII Detector", version="1.0.0")
# Load model at startup — stays in memory
ner = pipeline(
"ner",
model="Ramosh99/mixcode_pii_detector",
aggregation_strategy="average",
token=os.getenv("HF_TOKEN"),
)
class PredictRequest(BaseModel):
text: str
@app.get("/health")
def health():
return {"status": "ok", "model": "Ramosh99/mixcode_pii_detector"}
@app.post("/predict")
def predict(req: PredictRequest):
raw = ner(req.text)
entities = []
for item in raw:
label = item.get("entity_group", "")
if label.startswith(("B-", "I-")):
label = label[2:]
entities.append({
"span": item["word"].strip(),
"start": item["start"],
"end": item["end"],
"entity_group": label,
"confidence": round(float(item["score"]), 4),
})
return {"entities": entities}
|