Update app.py

app.py

@@ -211,7 +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.", }]
 ]
 
-
 @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:
@@ -240,14 +239,12 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
             elif choice == "doc":
                 message['text'] = contents + "\n\n\n" + "{} files uploaded.\n".format(1) + message['text']
                 conversation.append({"role": "user", "content": message['text']})
-                # format_msg = contents + "\n\n\n" + "{} files uploaded.\n".format(1) + message['text']
-                # conversation.append({"role": "user", "content": format_msg})
         else:
             conversation.append({"role": "user", "content": message['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)
+        # Convierte solo el contenido relevante a IDs
+        # 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)
 
         generate_kwargs = dict(
             max_length=max_length,
@@ -259,30 +256,92 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
             eos_token_id=[151329, 151336, 151338],
         )
 
-        gen_kwargs = {**input_ids, **generate_kwargs}
-
-        for entry in conversation:
-            print(f"Role: {entry['role']}, Content: {entry.get('content', '')}")
-
         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)
 
-        text_original = message['text'].strip()
-        generated_text_cleaned = generated_text.replace(text_original, "").strip()
-
         print(" ")
         print("---------")
         print("Text: ")
         print("  ")
-        print(generated_text_cleaned)
-
+        print(generated_text)
 
-        return PlainTextResponse(generated_text_cleaned)
+        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:
+#         model = AutoModelForCausalLM.from_pretrained(
+#             MODEL_ID,
+#             torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+#             low_cpu_mem_usage=True,
+#             trust_remote_code=True
+#         )
+
+#         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"]
+
+#             with open(file_name, "wb") as f:
+#                 f.write(file_content.read())
+
+#             choice, contents = mode_load(file_name)
+
+#             if choice == "image":
+#                 conversation.append({"role": "user", "image": contents, "content": 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']})
+#                 # format_msg = contents + "\n\n\n" + "{} files uploaded.\n".format(1) + message['text']
+#                 # conversation.append({"role": "user", "content": format_msg})
+#         else:
+#             conversation.append({"role": "user", "content": message['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,
+#             do_sample=True,
+#             top_p=top_p,
+#             top_k=top_k,
+#             temperature=temperature,
+#             repetition_penalty=penalty,
+#             eos_token_id=[151329, 151336, 151338],
+#         )
+
+#         gen_kwargs = {**input_ids, **generate_kwargs}
+
+#         for entry in conversation:
+#             print(f"Role: {entry['role']}, Content: {entry.get('content', '')}")
+
+#         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)
+
+#         text_original = message['text'].strip()
+#         generated_text_cleaned = generated_text.replace(text_original, "").strip()
+
+#         print(" ")
+#         print("---------")
+#         print("Text: ")
+#         print("  ")
+#         print(generated_text_cleaned)
+
+
+#         return PlainTextResponse(generated_text_cleaned)
+#     except Exception as e:
+#         return PlainTextResponse(f"Error: {str(e)}")
+
+
         
 # 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: