Update app.py

app.py

@@ -1,49 +1,42 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
-import numpy as np
+from transformers import pipeline
 
-app = FastAPI()
+def get_classifier():
+    classifier = pipeline(
+        "zero-shot-classification",
+        model="sberbank-ai/rugpt3small_based_on_gpt2",
+        framework="pt"
+    )
+    return classifier
 
-# Load model and tokenizer
-MODEL_NAME = "xlm-roberta-base"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=3)
-
-class EmailRequest(BaseModel):
-    text: str
-
-class EmailResponse(BaseModel):
-    category: int
-    confidence: float
-
-LABELS = {
-    0: "Клиент хочет назначить встречу",
-    1: "Клиент не заинтересован / нет времени / отказывается",
-    2: "Клиент задаёт уточняющие вопросы"
-}
-
-@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)
+def classify_email(text):
+    classifier = get_classifier()
+    
+    candidate_labels = [
+        "Клиент хочет назначить встречу",
+        "Клиент не заинтересован или отказывается",
+        "Клиент задаёт уточняющие вопросы"
+    ]
+    
+    result = classifier(
+        text,
+        candidate_labels,
+        hypothesis_template="Это письмо о том, что {}."
+    )
+    
+    # Получаем индекс наиболее вероятной метки (0, 1 или 2)
+    label_index = result["labels"].index(result["labels"][0])
     
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+    # Возвращаем категорию (1, 2 или 3) и уверенность
+    return {
+        "category": label_index + 1,
+        "confidence": result["scores"][label_index],
+        "label": result["labels"][0]
+    }
 
+# Пример использования
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    test_text = "Добрый день! Можно ли узнать подробнее о ваших услугах и ценах?"
+    result = classify_email(test_text)
+    print(f"Категория: {result['category']}")
+    print(f"Уверенность: {result['confidence']:.2f}")
+    print(f"Метка: {result['label']}")