Update app.py

app.py

@@ -211,7 +211,7 @@ EXAMPLES = [
 
 # Definir la función simple_chat
 @spaces.GPU()
-def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096, top_p: float = 1, top_k: int = 10, penalty: float = 1.0):
+async def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096, top_p: float = 1, top_k: int = 10, penalty: float = 1.0):
     # Cargar el modelo preentrenado
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
@@ -219,7 +219,7 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
         low_cpu_mem_usage=True,
         trust_remote_code=True
     )
-        
+
     conversation = []
 
     if "file" in message and message["file"]:
@@ -236,12 +236,10 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
     # Preparar entrada para el modelo
     input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True,
                                               return_tensors="pt", return_dict=True).to(model.device)
-    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
-
+    
     # Configurar parámetros de generación
     generate_kwargs = dict(
         max_length=max_length,
-        streamer=streamer,
         do_sample=True,
         top_p=top_p,
         top_k=top_k,
@@ -249,19 +247,13 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
         repetition_penalty=penalty,
         eos_token_id=[151329, 151336, 151338],
     )
-    gen_kwargs = {**input_ids, **generate_kwargs}
 
-    # Generar respuesta de manera asíncrona
-    def generate():
-        with torch.no_grad():
-            thread = Thread(target=model.generate, kwargs=gen_kwargs)
-            thread.start()
-            buffer = ""
-            for new_text in streamer:
-                buffer += new_text
-                yield buffer.encode('utf-8')
-
-    #return StreamingResponse(generate(), media_type="text/plain")
+    # Generar respuesta
+    with torch.no_grad():
+        generated_ids = model.generate(input_ids['input_ids'], **generate_kwargs)
+        generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+
+    # Devolver la respuesta completa
     return PlainTextResponse(generated_text)
 
 @app.post("/chat/")