Create classifier.py

classifier.py

@@ -0,0 +1,46 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from typing import List
+import torch
+from transformers import AutoTokenizer, XLMRobertaForSequenceClassification
+
+# === Конфигурация ===
+MODEL_NAME = "xlm-roberta-large"
+
+CATEGORIES = [
+    "politique", "woke", "racism", "crime",
+    "police_abuse", "corruption", "hate_speech", "activism"
+]
+
+# === Загрузка модели ===
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = XLMRobertaForSequenceClassification.from_pretrained(
+    MODEL_NAME,
+    num_labels=len(CATEGORIES)
+)
+model.eval()
+
+# === FastAPI приложение ===
+app = FastAPI(title="Multilabel Text Classifier API")
+
+# === Схема запроса ===
+class TextRequest(BaseModel):
+    text: str
+
+# === Логика классификации ===
+def classify_message(message: str) -> List[str]:
+    inputs = tokenizer(message, return_tensors="pt", truncation=True)
+    with torch.no_grad():
+        logits = model(**inputs).logits
+
+    probs = torch.sigmoid(logits)[0]
+    selected = [CATEGORIES[i] for i, p in enumerate(probs) if p > 0.5]
+    return selected or ["neutral"]
+
+# === Эндпоинт ===
+@app.post("/classify")
+def classify(request: TextRequest):
+    categories = classify_message(request.text)
+    return {
+        "categories": categories
+    }