Update app.py

app.py

@@ -211,77 +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:
-        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)
-
-        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":
-                #input_text = message['text']
-                conversation.append({"role": "user", "image": contents, "content": message['text']})
-            elif choice == "doc":
-                #input_text = contents + "\n\n\n" + "{} files uploaded.\n".format(1) + message['text']
-                format_msg = contents + "\n\n\n" + "{} files uploaded.\n" + message['text']
-                conversation.append({"role": "user", "content": format_msg})
-        else:
-            #input_text = message['text']
-            contents = None
-            conversation.append({"role": "user", "content": message['text']})
-
-        # conversation = [{"role": "user", "content": input_text}]
-        input_ids = tokenizer(conversation[-1]['content'], return_tensors="pt").to(model.device)
-
-        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]
-        )
-
-        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)
-
-        # Process to remove any prefix or unwanted prompt
-        # text_original = input_text.strip()
-        # results_text = generated_text[len(text_original):].strip()
-
-        text_original = message['text'].strip()
-        results_text = generated_text.replace(text_original, "").strip()
-
-        print("  ")
-        print("------")
-        print("  ")
-        print(generated_text)
-        print("  ")
-        print("------")
-
-        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:
@@ -294,8 +223,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"]
@@ -306,18 +233,19 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
 #             choice, contents = mode_load(file_name)
 
 #             if choice == "image":
+#                 #input_text = message['text']
 #                 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})
+#                 #input_text = contents + "\n\n\n" + "{} files uploaded.\n".format(1) + message['text']
+#                 format_msg = contents + "\n\n\n" + "{} files uploaded.\n" + message['text']
+#                 conversation.append({"role": "user", "content": format_msg})
 #         else:
+#             #input_text = message['text']
+#             contents = None
 #             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)
+#         # conversation = [{"role": "user", "content": input_text}]
+#         input_ids = tokenizer(conversation[-1]['content'], return_tensors="pt").to(model.device)
 
 #         generate_kwargs = dict(
 #             max_length=max_length,
@@ -326,34 +254,106 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
 #             top_k=top_k,
 #             temperature=temperature,
 #             repetition_penalty=penalty,
-#             eos_token_id=[151329, 151336, 151338],
+#             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)
 
+#         # Process to remove any prefix or unwanted prompt
+#         # text_original = input_text.strip()
+#         # results_text = generated_text[len(text_original):].strip()
+
 #         text_original = message['text'].strip()
-#         generated_text_cleaned = generated_text.replace(text_original, "").strip()
+#         results_text = generated_text.replace(text_original, "").strip()
 
-#         print(" ")
-#         print("---------")
-#         print("Text: ")
 #         print("  ")
-#         print(generated_text_cleaned)
-
+#         print("------")
+#         print("  ")
+#         print(generated_text)
+#         print("  ")
+#         print("------")
 
-#         return PlainTextResponse(generated_text_cleaned)
+#         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:
+        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)}")
+
+
+
 @app.post("/chat/")
 async def test_endpoint(
     text: str = Form(...),