Update app.py

app.py

@@ -1,16 +1,21 @@
-import gradio as gr
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import os
+import uvicorn
+import threading
 
-# Reduzir verbosidade dos warnings
+# Configurações
 os.environ["TRANSFORMERS_VERBOSITY"] = "error"
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-# TinyLlama - modelo leve e eficiente
+# Modelo
 MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
-print("Carregando TinyLlama 1.1B...")
+print("🦙 Carregando TinyLlama para API...")
 
+# Carregar modelo
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
@@ -19,95 +24,142 @@ model = AutoModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True
 )
 
-# Configurar pad token
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
-print("✅ Modelo carregado! Interface iniciando...")
+print("✅ Modelo carregado! API iniciando...")
 
-def chat_response(message, max_tokens, temperature):
-    """Função principal de chat"""
+# FastAPI app
+app = FastAPI(
+    title="TinyLlama Chat API",
+    description="API REST para TinyLlama 1.1B",
+    version="1.0.0"
+)
+
+# Modelos Pydantic
+class ChatRequest(BaseModel):
+    message: str
+    max_tokens: int = 200
+    temperature: float = 0.7
+
+class ChatResponse(BaseModel):
+    response: str
+    status: str = "success"
+
+# Lock para thread safety
+model_lock = threading.Lock()
+
+def generate_response(message: str, max_tokens: int = 200, temperature: float = 0.7) -> str:
+    """Gerar resposta com o modelo"""
     try:
-        # Template do TinyLlama
-        prompt = f"<|system|>\nVocê é um assistente útil. Responda de forma clara e concisa.<|user|>\n{message}<|assistant|>\n"
-        
-        # Tokenizar
-        inputs = tokenizer(
-            prompt, 
-            return_tensors="pt", 
-            truncation=True, 
-            max_length=1200,
-            padding=False
-        )
-        
-        # Gerar resposta (sem early_stopping para evitar warning)
-        with torch.no_grad():
-            outputs = model.generate(
-                inputs.input_ids,
-                attention_mask=inputs.attention_mask,
-                max_new_tokens=max_tokens,
-                temperature=temperature,
-                do_sample=True,
-                top_p=0.9,
-                repetition_penalty=1.1,
-                pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=tokenizer.eos_token_id
+        with model_lock:
+            prompt = f"<|system|>\nVocê é um assistente útil. Responda de forma clara e concisa.<|user|>\n{message}<|assistant|>\n"
+            
+            inputs = tokenizer(
+                prompt,
+                return_tensors="pt",
+                truncation=True,
+                max_length=1000,
+                padding=False
+            )
+            
+            with torch.no_grad():
+                outputs = model.generate(
+                    inputs.input_ids,
+                    max_new_tokens=min(max_tokens, 300),
+                    temperature=max(0.1, min(temperature, 1.0)),
+                    do_sample=True,
+                    top_p=0.9,
+                    repetition_penalty=1.1,
+                    pad_token_id=tokenizer.eos_token_id,
+                    eos_token_id=tokenizer.eos_token_id
+                )
+            
+            response = tokenizer.decode(
+                outputs[0][len(inputs.input_ids[0]):],
+                skip_special_tokens=True
             )
+            
+            # Limpar resposta
+            response = response.split("<|user|>")[0].split("<|system|>")[0].strip()
+            
+            return response if response else "Não consegui gerar uma resposta."
+            
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Erro na geração: {str(e)}")
+
+# Endpoints da API
+
+@app.get("/")
+async def root():
+    """Endpoint raiz - informações da API"""
+    return {
+        "message": "TinyLlama Chat API",
+        "model": MODEL_NAME,
+        "endpoints": {
+            "POST /chat": "Enviar mensagem para o modelo",
+            "GET /health": "Verificar status da API",
+            "GET /docs": "Documentação interativa"
+        }
+    }
+
+@app.get("/health")
+async def health_check():
+    """Verificar se a API está funcionando"""
+    return {
+        "status": "healthy",
+        "model_loaded": True,
+        "model_name": MODEL_NAME
+    }
+
+@app.post("/chat", response_model=ChatResponse)
+async def chat_endpoint(request: ChatRequest):
+    """Endpoint principal para chat"""
+    if not request.message or not request.message.strip():
+        raise HTTPException(status_code=400, detail="Mensagem não pode estar vazia")
+    
+    try:
+        response = generate_response(
+            message=request.message,
+            max_tokens=request.max_tokens,
+            temperature=request.temperature
+        )
         
-        # Extrair resposta
-        new_tokens = outputs[0][len(inputs.input_ids[0]):]
-        response = tokenizer.decode(new_tokens, skip_special_tokens=True)
+        return ChatResponse(response=response)
         
-        # Limpar resposta
-        response = response.split("<|user|>")[0]
-        response = response.split("<|system|>")[0]
-        response = response.strip()
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/chat")
+async def chat_get(message: str, max_tokens: int = 200, temperature: float = 0.7):
+    """Endpoint GET para chat (mais simples de testar)"""
+    if not message or not message.strip():
+        raise HTTPException(status_code=400, detail="Parâmetro 'message' é obrigatório")
+    
+    try:
+        response = generate_response(
+            message=message,
+            max_tokens=max_tokens,
+            temperature=temperature
+        )
         
-        return response if response else "Não consegui gerar uma resposta. Tente reformular sua pergunta."
+        return {"response": response, "status": "success"}
         
     except Exception as e:
-        return f"Erro: {str(e)}"
-
-# Interface Gradio simples e funcional
-interface = gr.Interface(
-    fn=chat_response,
-    inputs=[
-        gr.Textbox(
-            label="💬 Sua pergunta",
-            placeholder="Digite sua pergunta aqui...",
-            lines=2
-        ),
-        gr.Slider(
-            minimum=50,
-            maximum=400,
-            value=200,
-            step=10,
-            label="🔢 Tokens máximos"
-        ),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.2,
-            value=0.7,
-            step=0.1,
-            label="🌡️ Criatividade"
-        )
-    ],
-    outputs=gr.Textbox(
-        label="🤖 Resposta do TinyLlama",
-        lines=5
-    ),
-    title="🦙 TinyLlama Chat API",
-    description="Modelo de IA leve (2.2GB) otimizado para Hugging Face Spaces gratuito",
-    theme="default",
-    # Sem examples para evitar cache/erros
-    allow_flagging="never"
-)
+        raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
-    print("🚀 Iniciando servidor...")
-    interface.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        show_error=False
+    print("🚀 Iniciando servidor FastAPI...")
+    print("📡 API estará disponível em:")
+    print("   - GET  /           (informações)")
+    print("   - GET  /health     (status)")
+    print("   - POST /chat       (principal)")
+    print("   - GET  /chat       (teste simples)")
+    print("   - GET  /docs       (documentação)")
+    
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=7860,
+        log_level="error"  # Reduzir logs
     )