Update app.py

app.py

@@ -1,41 +1,49 @@
-from transformers import pipeline
-from flask import Flask, request, jsonify
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+import numpy as np
 
-app = Flask(__name__)
+app = FastAPI()
 
-# Загружаем zero-shot классификатор
-classifier = pipeline("zero-shot-classification", model="typeform/distilbert-base-uncased-mnli")
+# Load model and tokenizer
+MODEL_NAME = "xlm-roberta-base"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=3)
 
-# Гипотезы
-LABELS = {
-    "wants_meeting": "Клиент хочет назначить встречу или обсудить время",
-    "not_interested": "Клиент не заинтересован во встрече или у него нет времени",
-    "asking_questions": "Клиент задает уточняющие вопросы по теме"
-}
-
-@app.route("/analyze", methods=["POST"])
-def analyze():
-    data = request.get_json()
-    emails = data.get("emails", [])
+class EmailRequest(BaseModel):
+    text: str
 
-    if not emails or not isinstance(emails, list):
-        return jsonify({"error": "Field 'emails' must be a non-empty list"}), 400
+class EmailResponse(BaseModel):
+    category: int
+    confidence: float
 
-    results = []
-    for email in emails:
-        prediction = classifier(email, list(LABELS.values()), multi_label=True)
-        scored_labels = dict(zip(prediction["labels"], prediction["scores"]))
-
-        # Вывод в формате {label: score}
-        result = {
-            "text": email,
-            "intents": {
-                key: round(scored_labels[label], 4) for key, label in LABELS.items()
-            }
-        }
-        results.append(result)
+LABELS = {
+    0: "Клиент хочет назначить встречу",
+    1: "Клиент не заинтересован / нет времени / отказывается",
+    2: "Клиент задаёт уточняющие вопросы"
+}
 
-    return jsonify(results)
+@app.post("/classify", response_model=EmailResponse)
+async def classify_email(request: EmailRequest):
+    try:
+        # Tokenize the input text
+        inputs = tokenizer(request.text, return_tensors="pt", truncation=True, max_length=512)
+        
+        # Get model predictions
+        with torch.no_grad():
+            outputs = model(**inputs)
+            predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
+        
+        # Get the predicted class and confidence
+        predicted_class = torch.argmax(predictions).item()
+        confidence = predictions[0][predicted_class].item()
+        
+        return EmailResponse(category=predicted_class + 1, confidence=confidence)
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)