Update app.py

app.py

@@ -211,10 +211,6 @@ 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:
@@ -244,7 +240,7 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
             input_text = message['text']
 
         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)
+        input_ids = tokenizer(conversation[-1]['content'], return_tensors="pt").to(model.device)
         streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
 
         generate_kwargs = dict(
@@ -275,19 +271,25 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
         # Wait for the generation thread to finish
         thread.join()
 
+        # Process to remove any prefix or unwanted prompt
+        text_original = input_text.strip()
+        results_text = buffer[len(text_original):].strip()
+
         print("  ")
         print("------")
         print("  ")
-        print(buffer)
+        print(results_text)
         print("  ")
         print("------")
 
-        return PlainTextResponse(buffer)
+        return PlainTextResponse(results_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: