Update app.py

app.py

@@ -1,119 +1,71 @@
-import torch
-from fastapi import FastAPI
-from pydantic import BaseModel
-from typing import List, Optional, Union
-# !!! ИСПРАВЛЕНИЕ: Используем правильный класс для Seq2Seq моделей, как T5/FRIDA !!!
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, BitsAndBytesConfig
-
-# --- 1. Конфигурация и загрузка модели ---
-
-# Имя модели на Hugging Face
-MODEL_NAME = "ai-forever/FRIDA"
-# Имя, которое будет возвращаться в ответах API (может быть любым)
-MODEL_ALIAS = "frida-v1"
-
-print("Starting model loading process...")
-
-# Конфигурация квантизации для экономии памяти (ОБЯЗАТЕЛЬНО для бесплатного тарифа)
-quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+from typing import List, Optional
+import numpy as np
+
+from sentence_transformers import SentenceTransformer
+
+app = FastAPI(title="FRIDA Embedding API", version="1.0")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
-# Загружаем токенизатор
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
-
-# !!! ИСПРАВЛЕНИЕ: Загружаем модель с помощью AutoModelForSeq2SeqLM !!!
-model = AutoModelForSeq2SeqLM.from_pretrained(
-    MODEL_NAME,
-    quantization_config=quantization_config,
-    device_map="auto",  # Автоматически распределит модель по доступным ресурсам
-    trust_remote_code=True
-)
-model.eval() # Переводим модель в режим инференса
-
-print("Model and tokenizer loaded successfully.")
-
-# --- 2. Определение Pydantic-моделей для имитации OpenAI API ---
-
-class ChatMessage(BaseModel):
-    role: str
-    content: str
-
-class ChatCompletionRequest(BaseModel):
-    model: str # Будет проигнорировано, но нужно для совместимости
-    messages: List[ChatMessage]
-    temperature: Optional[float] = 0.7
-    max_tokens: Optional[int] = 1024
-
-class ChatCompletionChoice(BaseModel):
-    index: int
-    message: ChatMessage
-    finish_reason: str = "stop"
-
-class UsageInfo(BaseModel):
-    prompt_tokens: int
-    completion_tokens: int
-    total_tokens: int
-
-class ChatCompletionResponse(BaseModel):
-    id: str = "chatcmpl-mock"
-    object: str = "chat.completion"
-    created: int = 0
-    model: str = MODEL_ALIAS
-    choices: List[ChatCompletionChoice]
-    usage: UsageInfo
-
-# --- 3. Создание FastAPI приложения ---
-app = FastAPI()
-
-# --- 4. Реализация эндпоинта /v1/chat/completions ---
-
-@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
-async def create_chat_completion(request: ChatCompletionRequest):
-    """
-    Основная функция, которая принимает запрос и генерирует ответ.
-    """
-    print(f"Received request: {request.dict()}")
-    
-    # Преобразуем сообщения из формата OpenAI в единую строку для T5 модели
-    prompt_text = "\n".join([f"{msg.role}: {msg.content}" for msg in request.messages])
-
-    print(f"Formatted prompt for FRIDA:\n{prompt_text}")
-
-    # Кодируем текст в токены
-    inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device)
-    prompt_tokens_count = inputs["input_ids"].shape[1]
-
-    # Генерируем ответ от модели
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=request.max_tokens,
-        temperature=request.temperature,
-        do_sample=True,
-        eos_token_id=tokenizer.eos_token_id,
-        pad_token_id=tokenizer.pad_token_id
-    )
-
-    # Декодируем сгенерированные токены обратно в текст
-    response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Считаем токены ответа
-    completion_tokens_count = outputs[0].shape[0]
-
-    print(f"Generated response: {response_text}")
-
-    # Формируем ответ в формате OpenAI
-    response_message = ChatMessage(role="assistant", content=response_text)
-    choice = ChatCompletionChoice(index=0, message=response_message)
-    usage = UsageInfo(
-        prompt_tokens=prompt_tokens_count,
-        completion_tokens=completion_tokens_count,
-        total_tokens=prompt_tokens_count + completion_tokens_count
-    )
-
-    return ChatCompletionResponse(choices=[choice], usage=usage)
-
-@app.get("/")
-def health_check():
-    return {"status": "ok", "model_name": MODEL_ALIAS}
+MODEL_NAME = "ai-forever/FRIDA"
+model = SentenceTransformer(MODEL_NAME)
+EMBED_DIM = model.get_sentence_embedding_dimension()
+
+SUPPORTED_PROMPTS = [
+    "search_query",
+    "search_document",
+    "paraphrase",
+    "categorize",
+    "categorize_sentiment",
+    "categorize_topic",
+    "categorize_entailment",
+]
+
+class EmbedRequest(BaseModel):
+    texts: List[str] = Field(..., description="Список текстов")
+    prompt_name: Optional[str] = Field("search_document", description="FRIDA prompt_name")
+
+class EmbedResponse(BaseModel):
+    embeddings: List[List[float]]
+    dim: int
+
+@app.get("/health")
+def health():
+    return {"status": "ok"}
+
+@app.get("/metadata")
+def metadata():
+    return {
+        "model": MODEL_NAME,
+        "embedding_dim": EMBED_DIM,
+        "pooling": "cls",
+        "prompts_supported": SUPPORTED_PROMPTS,
+    }
+
+@app.post("/embed", response_model=EmbedResponse)
+def embed(req: EmbedRequest):
+    if not req.texts:
+        raise HTTPException(status_code=400, detail="texts must be non-empty")
+    prompt = req.prompt_name or "search_document"
+    if prompt not in SUPPORTED_PROMPTS:
+        raise HTTPException(status_code=400, detail=f"Unsupported prompt_name: {prompt}")
+
+    vectors = model.encode(
+        req.texts,
+        convert_to_numpy=True,
+        prompt_name=prompt,
+        normalize_embeddings=True,
+        batch_size=min(16, max(1, len(req.texts))),
+        show_progress_bar=False,
+    ).astype(np.float32)
+
+    return {"embeddings": vectors.tolist(), "dim": int(vectors.shape[1])}