Update app/main.py

app/main.py

@@ -3,65 +3,61 @@ from typing import List, Optional
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
-import ctranslate2
-from transformers import AutoTokenizer
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
 logger = logging.getLogger(__name__)
 
-# ─── Config via ENV ───────────────────────────────────────────────────────────
-CT2_MODEL_DIR  = os.getenv("CT2_MODEL_DIR",  "/app/ct2_model")
-TOKENIZER_DIR  = os.getenv("TOKENIZER_DIR",  "/app/tokenizer")
-CT2_MODEL_ID   = os.getenv("CT2_MODEL_ID",   "limcheekin/flan-t5-xxl-ct2")
-TOKENIZER_ID   = os.getenv("TOKENIZER_ID",   "google/flan-t5-xxl")
+MODEL_ID       = os.getenv("MODEL_ID",       "google/flan-t5-large")
 MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", "1024"))
-MAX_INPUT_LEN  = int(os.getenv("MAX_INPUT_LEN",  "512"))   # 16384 se usar long-t5
-INTER_THREADS  = int(os.getenv("INTER_THREADS",  "2"))
-INTRA_THREADS  = int(os.getenv("INTRA_THREADS",  "2"))
+MAX_INPUT_LEN  = int(os.getenv("MAX_INPUT_LEN",  "512"))
 
-# ─── App ─────────────────────────────────────────────────────────────────────
-app = FastAPI(title="T2T OpenAI-Compatible API", version="2.0.0")
+app = FastAPI(title="T2T OpenAI-Compatible API", version="3.0.0")
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"], allow_methods=["*"], allow_headers=["*"]
+    allow_origins=["*"], allow_methods=["*"], allow_headers=["*"],
 )
 
-_translator: ctranslate2.Translator = None
-_tokenizer:  AutoTokenizer          = None
+_pipe      = None
+_tokenizer = None
 
-# ─── Startup ─────────────────────────────────────────────────────────────────
 @app.on_event("startup")
 def load_model():
-    global _translator, _tokenizer
-    logger.info(f"⏳ Carregando CT2 model: {CT2_MODEL_DIR}")
-    _translator = ctranslate2.Translator(
-        CT2_MODEL_DIR,
-        device        = "cpu",
-        compute_type  = "int8",
-        inter_threads = INTER_THREADS,
-        intra_threads = INTRA_THREADS,
+    global _pipe, _tokenizer
+    logger.info(f"⏳ Carregando {MODEL_ID} …")
+    _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False)
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype  = torch.float32,
+        low_cpu_mem_usage = True,
     )
-    logger.info(f"⏳ Carregando tokenizer: {TOKENIZER_DIR}")
-    _tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_DIR)
-    logger.info("✅ Model pronto!")
+    model.eval()
+    _pipe = pipeline(
+        "text2text-generation",
+        model     = model,
+        tokenizer = _tokenizer,
+        device    = -1,          # força CPU
+    )
+    logger.info(f"✅ {MODEL_ID} pronto!")
 
-# ─── Schemas (espelho exato do OpenAI) ───────────────────────────────────────
+# ── Schemas (OpenAI-compatible) ───────────────────────────────────────────
 class Message(BaseModel):
-    role: str      # "system" | "user" | "assistant"
+    role: str
     content: str
 
 class ResponseFormat(BaseModel):
     type: str = "text"
 
 class ChatCompletionRequest(BaseModel):
-    model: str                               = Field(default=CT2_MODEL_ID)
+    model: str                              = Field(default=MODEL_ID)
     messages: List[Message]
-    temperature: float                       = 0.7
-    top_p: float                             = 0.9
-    max_completion_tokens: Optional[int]     = None
+    temperature: float                      = 0.7
+    top_p: float                            = 0.9
+    max_completion_tokens: Optional[int]    = None
     max_tokens: Optional[int]               = None
     response_format: Optional[ResponseFormat] = None
-    stream: bool                             = False
+    stream: bool                            = False
 
     class Config:
         populate_by_name = True
@@ -71,29 +67,25 @@ class ChoiceMessage(BaseModel):
     content: str
 
 class Choice(BaseModel):
-    index:         int
-    message:       ChoiceMessage
+    index: int
+    message: ChoiceMessage
     finish_reason: str = "stop"
 
 class Usage(BaseModel):
-    prompt_tokens:     int
+    prompt_tokens: int
     completion_tokens: int
-    total_tokens:      int
+    total_tokens: int
 
 class ChatCompletionResponse(BaseModel):
-    id:      str
-    object:  str  = "chat.completion"
+    id: str
+    object: str  = "chat.completion"
     created: int
-    model:   str
+    model: str
     choices: List[Choice]
-    usage:   Usage
+    usage: Usage
 
-# ─── Helpers ─────────────────────────────────────────────────────────────────
+# ── Helpers ───────────────────────────────────────────────────────────────
 def messages_to_prompt(messages: List[Message]) -> str:
-    """
-    Converte lista de mensagens em prompt único para modelos seq2seq.
-    Preserva contexto system + histórico de conversa.
-    """
     parts = []
     for m in messages:
         if m.role == "system":
@@ -102,95 +94,61 @@ def messages_to_prompt(messages: List[Message]) -> str:
             parts.append(f"User: {m.content}")
         elif m.role == "assistant":
             parts.append(f"Assistant: {m.content}")
-    return "\n".join(parts)
+    return " ".join(parts)
 
-def count_tokens(text: str) -> int:
+def token_count(text: str) -> int:
     return len(_tokenizer(text, add_special_tokens=False)["input_ids"])
 
-# ─── Endpoints ───────────────────────────────────────────────────────────────
+# ── Endpoints ─────────────────────────────────────────────────────────────
 @app.get("/")
 def root():
-    return {
-        "status": "ok",
-        "model": CT2_MODEL_ID,
-        "max_input_tokens": MAX_INPUT_LEN,
-        "max_output_tokens": MAX_NEW_TOKENS,
-    }
+    return {"status": "ok", "model": MODEL_ID}
 
 @app.get("/health")
 def health():
-    return {
-        "status":       "healthy",
-        "model":        CT2_MODEL_ID,
-        "model_loaded": _translator is not None,
-    }
+    return {"status": "healthy", "model": MODEL_ID, "ready": _pipe is not None}
 
 @app.get("/v1/models")
 def list_models():
-    return {
-        "object": "list",
-        "data": [{
-            "id":       CT2_MODEL_ID,
-            "object":   "model",
-            "owned_by": "huggingface",
-        }],
-    }
+    return {"object": "list", "data": [
+        {"id": MODEL_ID, "object": "model", "owned_by": "huggingface"}
+    ]}
 
 @app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
 def chat_completions(req: ChatCompletionRequest):
     if req.stream:
-        raise HTTPException(501, "Streaming não suportado ainda.")
-
-    if _translator is None or _tokenizer is None:
-        raise HTTPException(503, "Modelo ainda não carregado.")
+        raise HTTPException(501, "Streaming não suportado.")
+    if _pipe is None:
+        raise HTTPException(503, "Modelo não carregado.")
 
     max_tokens = req.max_completion_tokens or req.max_tokens or MAX_NEW_TOKENS
     prompt     = messages_to_prompt(req.messages)
-
-    # Tokeniza com truncation para respeitar janela do modelo
-    encoded      = _tokenizer(
-        prompt,
-        return_tensors   = None,
-        truncation       = True,
-        max_length       = MAX_INPUT_LEN,
-        add_special_tokens = True,
-    )
-    input_tokens = [_tokenizer.convert_ids_to_tokens(encoded["input_ids"])]
-
-    do_sample = req.temperature > 0.05
+    do_sample  = req.temperature > 0.05
 
     try:
-        results = _translator.translate_batch(
-            input_tokens,
-            max_decoding_length  = max_tokens,
-            min_decoding_length  = 1,
-            beam_size            = 1 if do_sample else 4,
-            sampling_temperature = float(req.temperature) if do_sample else 1.0,
-            sampling_topk        = 50 if do_sample else 1,
-            sampling_topp        = float(req.top_p) if do_sample else 1.0,
-            repetition_penalty   = 1.2,   # evita repetições em textos longos
+        output = _pipe(
+            prompt,
+            max_new_tokens     = max_tokens,
+            truncation         = True,
+            temperature        = float(req.temperature) if do_sample else 1.0,
+            top_p              = float(req.top_p)        if do_sample else 1.0,
+            do_sample          = do_sample,
+            repetition_penalty = 1.2,
         )
     except Exception as e:
         logger.error(f"Inference error: {e}")
-        raise HTTPException(500, f"Erro na inferência: {e}")
-
-    output_tokens = results[0].hypotheses[0]
-    generated     = _tokenizer.convert_tokens_to_string(output_tokens).strip()
+        raise HTTPException(500, str(e))
 
-    p_tok = count_tokens(prompt)
-    c_tok = count_tokens(generated)
+    text  = output[0]["generated_text"].strip()
+    p_tok = token_count(prompt)
+    c_tok = token_count(text)
 
     return ChatCompletionResponse(
         id      = f"chatcmpl-{uuid.uuid4().hex[:12]}",
         created = int(time.time()),
         model   = req.model,
-        choices = [
-            Choice(
-                index   = 0,
-                message = ChoiceMessage(content=generated),
-            )
-        ],
-        usage = Usage(
+        choices = [Choice(index=0, message=ChoiceMessage(content=text))],
+        usage   = Usage(
             prompt_tokens     = p_tok,
             completion_tokens = c_tok,
             total_tokens      = p_tok + c_tok,