Update app.py

app.py

@@ -117,37 +117,6 @@ def mode_load(path):
     else:
         raise gr.Error("Oops, unsupported files.")
 
-# def mode_load(file_obj):
-#     try:
-#         file_obj.seek(0)  # Asegúrate de que el puntero esté al inicio del archivo
-
-#         # Detecta el tipo de archivo basándote en los primeros bytes si es posible
-#         file_header = file_obj.read(4)
-#         file_obj.seek(0)  # Vuelve al inicio del archivo para procesamiento completo
-
-#         if file_header.startswith(b'%PDF'):
-#             content = extract_pdf(file_obj)
-#             choice = "doc"
-#         elif file_obj.name.endswith(".docx"):
-#             content = extract_docx(file_obj)
-#             choice = "doc"
-#         elif file_obj.name.endswith(".pptx"):
-#             content = extract_pptx(file_obj)
-#             choice = "doc"
-#         elif file_obj.name.endswith(".txt") or file_obj.name.endswith(".py") or file_obj.name.endswith(".json") or file_obj.name.endswith(".cpp") or file_obj.name.endswith(".md"):
-#             content = file_obj.read().decode('utf-8', errors='ignore')
-#             choice = "doc"
-#         elif file_obj.name.endswith((".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".webp")):
-#             content = Image.open(file_obj).convert('RGB')
-#             choice = "image"
-#         else:
-#             raise ValueError("Unsupported file type.")
-        
-#         return choice, content
-
-#     except Exception as e:
-#         raise ValueError(f"Error processing file: {str(e)}")
-
 
 @spaces.GPU()
 def stream_chat(message, history: list, temperature: float, max_length: int, top_p: float, top_k: int, penalty: float):
@@ -243,62 +212,6 @@ EXAMPLES = [
 ]
 
 
-# Definir la función simple_chat
-# @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,
-#             low_cpu_mem_usage=True,
-#             trust_remote_code=True
-#         )
-
-#         conversation = []
-
-#         if "file" in message and message["file"]:
-#             file_path = message["file"]
-#             choice, contents = mode_load(file_path)
-#             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:
-#             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)
-        
-#         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)
-
-#         return PlainTextResponse(generated_text)
-#     except Exception as e:
-#         return PlainTextResponse(f"Error: {str(e)}")
-
-# @app.post("/chat/")
-# async def test_endpoint(message: dict):
-#     if "text" not in message:
-#         raise HTTPException(status_code=400, detail="Missing 'text' in request body")
-
-#     if "file" not in message:
-#         print("Sin File")
-    
-#     response = simple_chat(message)
-#     return response
-
-
 @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:
@@ -342,6 +255,8 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
 
         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,
@@ -352,6 +267,23 @@ 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}
+
+        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("Buffer: ")
+        print("  ")
+        print(buffer)
+        print("  ")
+        print("--------------")
+
         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)
@@ -360,6 +292,67 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
         return PlainTextResponse(generated_text)
     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:
+#         model = AutoModelForCausalLM.from_pretrained(
+#             MODEL_ID,
+#             torch_dtype=torch.bfloat16,
+#             low_cpu_mem_usage=True,
+#             trust_remote_code=True
+#         )
+        
+#         tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+
+#         conversation = []
+
+#         # Acceder al contenido del archivo y al nombre del archivo
+#         if "file_content" in message and message["file_content"]:
+#             file_content = message["file_content"]
+#             file_name = message["file_name"]
+            
+#             # Guardar el archivo en un archivo temporal
+#             with open(file_name, "wb") as f:
+#                 f.write(file_content.read())
+            
+#             # Llamar a `mode_load` con el nombre del archivo
+#             choice, contents = mode_load(file_name)
+            
+#             if choice == "image":
+#                 conversation.append({"role": "user", "image": contents, "content": message['text']})
+#             elif choice == "doc":
+#                 format_msg = contents + "\n\n\n" + "{} files uploaded.\n".format(1) + message['text']
+#                 conversation.append({"role": "user", "content": format_msg})
+#         else:
+#             # Manejar caso donde no se sube archivo
+#             conversation.append({"role": "user", "content": message['text']})
+
+#         print("--------------")
+#         print("  ")
+#         print(conversation)
+#         print("  ")
+#         print("--------------")
+
+#         input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).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)
+
+
+#         return PlainTextResponse(generated_text)
+#     except Exception as e:
+#         return PlainTextResponse(f"Error: {str(e)}")