Update app.py

app.py

@@ -211,6 +211,10 @@ EXAMPLES = [
     [{"text": "Quiero armar un JSON, solo el JSON sin texto, que contenga los datos de la primera mitad de la tabla de la imagen (las primeras 10 jurisdicciones 901-910). Ten en cuenta que los valores numéricos son decimales de cuatro dígitos. La tabla contiene las siguientes columnas: Codigo, Nombre, Fecha Inicio, Fecha Cese, Coeficiente Ingresos, Coeficiente Gastos y Coeficiente Unificado. La tabla puede contener valores vacíos, en ese caso dejarlos como null. Cada fila de la tabla representa una jurisdicción con sus respectivos valores.", }]
 ]
 
+from threading import Thread
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import torch
+
 @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):
     try:
@@ -223,8 +227,6 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
 
         tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 
-        conversation = []
-
         if "file_content" in message and message["file_content"]:
             file_content = message["file_content"]
             file_name = message["file_name"]
@@ -235,14 +237,15 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
             choice, contents = mode_load(file_name)
 
             if choice == "image":
-                conversation.append({"role": "user", "image": contents, "content": message['text']})
+                input_text = message['text']
             elif choice == "doc":
-                message['text'] = contents + "\n\n\n" + "{} files uploaded.\n".format(1) + message['text']
-                conversation.append({"role": "user", "content": message['text']})
+                input_text = contents + "\n\n\n" + "{} files uploaded.\n".format(1) + message['text']
         else:
-            conversation.append({"role": "user", "content": message['text']})
+            input_text = message['text']
 
-        input_ids = tokenizer(conversation, return_tensors="pt").to(model.device)
+        conversation = [{"role": "user", "content": input_text}]
+        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)
 
         generate_kwargs = dict(
             max_length=max_length,
@@ -252,23 +255,35 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
             temperature=temperature,
             repetition_penalty=penalty,
             eos_token_id=[151329, 151336, 151338],
+            streamer=streamer
         )
 
-        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)
+        def generate_text():
+            with torch.no_grad():
+                model.generate(input_ids['input_ids'], **generate_kwargs)
+
+        # Run generation in a separate thread
+        thread = Thread(target=generate_text)
+        thread.start()
 
-        print(" ")
-        print("---------")
-        print("Text: ")
+        # Collect generated text in real-time
+        buffer = ""
+        for new_text in streamer:
+            buffer += new_text
+            yield buffer
+
+        print("  ")
+        print("------")
         print("  ")
-        print(generated_text)
+        print(buffer)
+        print("  ")
+        print("------")
 
-        return PlainTextResponse(generated_text)
     except Exception as e:
         return PlainTextResponse(f"Error: {str(e)}")
 
 
+
 # @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):
 #     try: