| """ |
| 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") |
|
|
| |
| 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} |
|
|
