Update app.py

app.py

@@ -1,23 +1,57 @@
-from fastapi import FastAPI
-from pydantic import BaseModel
+import gradio as gr
+import spaces
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-from agent import criar_arquivo
+modelo_id = "Qwen/Qwen2.5-Coder-7B-Instruct"
 
+tokenizer = AutoTokenizer.from_pretrained(modelo_id)
 
-app = FastAPI()
+modelo = None
 
 
-class Pedido(BaseModel):
-    tarefa:str
+@spaces.GPU
+def gerar_codigo(prompt):
+    global modelo
 
+    if modelo is None:
+        modelo = AutoModelForCausalLM.from_pretrained(
+            modelo_id,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
 
-@app.post("/criar")
-def criar(pedido:Pedido):
+    entradas = tokenizer(
+        prompt,
+        return_tensors="pt"
+    ).to(modelo.device)
 
-    arquivo = criar_arquivo(
-        pedido.tarefa
+    saida = modelo.generate(
+        **entradas,
+        max_new_tokens=512
     )
 
-    return {
-        "arquivo criado":arquivo
-    }
+    texto = tokenizer.decode(
+        saida[0],
+        skip_special_tokens=True
+    )
+
+    return texto
+
+
+app = gr.Interface(
+    fn=gerar_codigo,
+    inputs=gr.Textbox(
+        label="Pedido"
+    ),
+    outputs=gr.Textbox(
+        label="Resposta"
+    ),
+    title="Qwen2.5 Coder Agente"
+)
+
+
+app.launch(
+    server_name="0.0.0.0",
+    server_port=7860
+)