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']})
-
-
-            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)
+#     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']})
 
-    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)}")