Update app.py

app.py

@@ -1,7 +1,10 @@
-import gradio as gr
 import spaces
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+app = FastAPI()
 
 modelo_id = "Qwen/Qwen2.5-Coder-7B-Instruct"
 
@@ -10,8 +13,14 @@ tokenizer = AutoTokenizer.from_pretrained(modelo_id)
 modelo = None
 
 
+class Pedido(BaseModel):
+    model: str
+    messages: list
+
+
 @spaces.GPU
-def gerar_codigo(prompt):
+def gerar(prompt):
+
     global modelo
 
     if modelo is None:
@@ -21,37 +30,38 @@ def gerar_codigo(prompt):
             device_map="auto"
         )
 
-    entradas = tokenizer(
+    entrada = tokenizer(
         prompt,
         return_tensors="pt"
     ).to(modelo.device)
 
     saida = modelo.generate(
-        **entradas,
-        max_new_tokens=512
+        **entrada,
+        max_new_tokens=1024
     )
 
-    texto = tokenizer.decode(
+    return tokenizer.decode(
         saida[0],
         skip_special_tokens=True
     )
 
-    return texto
 
+@app.post("/v1/chat/completions")
+def chat(req: Pedido):
 
-app = gr.Interface(
-    fn=gerar_codigo,
-    inputs=gr.Textbox(
-        label="Pedido"
-    ),
-    outputs=gr.Textbox(
-        label="Resposta"
-    ),
-    title="Qwen2.5 Coder Agente"
-)
+    texto = req.messages[-1]["content"]
 
+    resposta = gerar(texto)
 
-app.launch(
-    server_name="0.0.0.0",
-    server_port=7860
-)
+    return {
+        "id": "qwen-local",
+        "object": "chat.completion",
+        "choices": [
+            {
+                "message": {
+                    "role": "assistant",
+                    "content": resposta
+                }
+            }
+        ]
+    }