Update app.py

app.py

@@ -211,157 +211,157 @@ 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, history: list, temperature: float = 0.8, max_length: int = 4096, top_p: float = 1, top_k: int = 10, penalty: float = 1.0):
+@spaces.GPU()
+def simple_chat(message, history: list, temperature: float = 0.8, max_length: int = 4096, top_p: float = 1, top_k: int = 10, penalty: float = 1.0):
     
-#     model = AutoModelForCausalLM.from_pretrained(
-#         MODEL_ID,
-#         torch_dtype=torch.bfloat16,
-#         low_cpu_mem_usage=True,
-#         trust_remote_code=True
-#     )
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=True,
+        trust_remote_code=True
+    )
         
-#     print(f'message is - {message}')
-#     print(f'history is - {history}')
-#     conversation = []
-#     prompt_files = []
-#     if message["files"]:
-#         choice, contents = mode_load(message["files"][-1])
-#         if choice == "image":
-#             conversation.append({"role": "user", "image": contents, "content": message['text']})
-#         elif choice == "doc":
-#             format_msg = contents + "\n\n\n" + "{} files uploaded.\n" + message['text']
-#             conversation.append({"role": "user", "content": format_msg})
-#     else:
-#         if len(history) == 0:
-#             # raise gr.Error("Please upload an image first.")
-#             contents = None
-#             conversation.append({"role": "user", "content": message['text']})
-#         else:
-#             # image = Image.open(history[0][0][0])
-#             for prompt, answer in history:
-#                 if answer is None:
-#                     prompt_files.append(prompt[0])
-#                     conversation.extend([{"role": "user", "content": ""}, {"role": "assistant", "content": ""}])
-#                 else:
-#                     conversation.extend([{"role": "user", "content": prompt}, {"role": "assistant", "content": answer}])
-#             if len(prompt_files) > 0:
-#                 choice, contents = mode_load(prompt_files[-1])
-#             else:
-#                 choice = ""
-#                 conversation.append({"role": "user", "image": "", "content": message['text']})
+    print(f'message is - {message}')
+    print(f'history is - {history}')
+    conversation = []
+    prompt_files = []
+    if message["files"]:
+        choice, contents = mode_load(message["files"][-1])
+        if choice == "image":
+            conversation.append({"role": "user", "image": contents, "content": message['text']})
+        elif choice == "doc":
+            format_msg = contents + "\n\n\n" + "{} files uploaded.\n" + message['text']
+            conversation.append({"role": "user", "content": format_msg})
+    else:
+        if len(history) == 0:
+            # raise gr.Error("Please upload an image first.")
+            contents = None
+            conversation.append({"role": "user", "content": message['text']})
+        else:
+            # image = Image.open(history[0][0][0])
+            for prompt, answer in history:
+                if answer is None:
+                    prompt_files.append(prompt[0])
+                    conversation.extend([{"role": "user", "content": ""}, {"role": "assistant", "content": ""}])
+                else:
+                    conversation.extend([{"role": "user", "content": prompt}, {"role": "assistant", "content": answer}])
+            if len(prompt_files) > 0:
+                choice, contents = mode_load(prompt_files[-1])
+            else:
+                choice = ""
+                conversation.append({"role": "user", "image": "", "content": message['text']})
 
 
-#             if choice == "image":
-#                 conversation.append({"role": "user", "image": contents, "content": message['text']})
-#             elif choice == "doc":
-#                 format_msg = contents + "\n\n\n" + "{} files uploaded.\n" + message['text']
-#                 conversation.append({"role": "user", "content": format_msg})
-#     print(f"Conversation is -\n{conversation}")
-
-#     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,
-#         streamer=streamer,
-#         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}
-
-#     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
-#         print(" ")
-#         print("---------")
-#         print("Text: ")
-#         print("  ")
-#         print(buffer)
-#         print(" ")
-#         print("---------")
+            if choice == "image":
+                conversation.append({"role": "user", "image": contents, "content": message['text']})
+            elif choice == "doc":
+                format_msg = contents + "\n\n\n" + "{} files uploaded.\n" + message['text']
+                conversation.append({"role": "user", "content": format_msg})
+    print(f"Conversation is -\n{conversation}")
+
+    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,
+        streamer=streamer,
+        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}
+
+    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
+        print(" ")
+        print("---------")
+        print("Text: ")
+        print("  ")
+        print(buffer)
+        print(" ")
+        print("---------")
 
 
 
 
-@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
-        )
+# @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)
+#         #tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 
-        conversation = []
+#         conversation = []
 
-        if "file_content" in message and message["file_content"]:
-            file_content = message["file_content"]
-            file_name = message["file_name"]
+#         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())
+#             with open(file_name, "wb") as f:
+#                 f.write(file_content.read())
 
-            choice, contents = mode_load(file_name)
+#             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']})
+#             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)
+#         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)
+#         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],
-        )
+#         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}
+#         gen_kwargs = {**input_ids, **generate_kwargs}
 
-        for entry in conversation:
-            print(f"Role: {entry['role']}, Content: {entry.get('content', '')}")
+#         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)
+#         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()
+#         text_original = message['text'].strip()
+#         generated_text_cleaned = generated_text.replace(text_original, "").strip()
 
-        print(" ")
-        print("---------")
-        print("Text: ")
-        print("  ")
-        print(generated_text_cleaned)
+#         print(" ")
+#         print("---------")
+#         print("Text: ")
+#         print("  ")
+#         print(generated_text_cleaned)
 
 
-        return PlainTextResponse(generated_text_cleaned)
-    except Exception as e:
-        return PlainTextResponse(f"Error: {str(e)}")
+#         return PlainTextResponse(generated_text_cleaned)
+#     except Exception as e:
+#         return PlainTextResponse(f"Error: {str(e)}")