server.py

import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import uvicorn
from typing import List, Optional
import time
import re
from functools import lru_cache

app = FastAPI(title="Spam Classifier")

MODEL_NAME = "super-apple/spam-classifier-ru"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"🔧 Используется устройство: {device}")

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
model.to(device)
model.eval()
print(f"✅ Модель {MODEL_NAME} загружена")

EXACT_BLOCK_PHRASES = []

SPAM_PATTERNS = []

def heuristic_check(text: str) -> bool:
    for phrase in EXACT_BLOCK_PHRASES:
        if phrase in text:
            return True
    for pattern in SPAM_PATTERNS:
        if re.search(pattern, text, re.IGNORECASE):
            return True
    return False

class ChatMessage(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    model: Optional[str] = None
    messages: List[ChatMessage]
    temperature: Optional[float] = 0.0
    max_tokens: Optional[int] = 1
    stream: Optional[bool] = False

class Choice(BaseModel):
    index: int
    message: ChatMessage

class ChatResponse(BaseModel):
    id: str = "spam-classifier"
    object: str = "chat.completion"
    created: int = 0
    model: str = MODEL_NAME
    choices: List[Choice]

@lru_cache(maxsize=1000)
def cached_classify(text: str) -> str:
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.softmax(outputs.logits, dim=1)
        spam_prob = probs[0][1].item()
        return "spam" if spam_prob > 0.5 else "ham"

@app.post("/v1/chat/completions", response_model=ChatResponse)
async def chat_completions(request: ChatRequest):
    user_messages = [m for m in request.messages if m.role == "user"]
    if not user_messages:
        raise HTTPException(status_code=400, detail="No user message found")
    
    text = user_messages[-1].content.strip()
    if not text:
        raise HTTPException(status_code=400, detail="Empty message")

    if heuristic_check(text):
        label = "spam"
        print(f"🔍 Эвристика определила спам: {text[:50]}...")
    else:
        label = cached_classify(text)
        print(f"🤖 Модель определила: {label} для: {text[:50]}...")

    response = ChatResponse(
        created=int(time.time()),
        choices=[
            Choice(
                index=0,
                message=ChatMessage(role="assistant", content=label)
            )
        ]
    )
    return response

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)