Update app.py

app.py

@@ -209,68 +209,68 @@ 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):
-#     # Cargar el modelo preentrenado
-#     model = AutoModelForCausalLM.from_pretrained(
-#         MODEL_ID,
-#         torch_dtype=torch.bfloat16,
-#         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):
+    # Cargar el modelo preentrenado
+    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"]})
-
-#     # Preparar entrada para el modelo
-#     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)
-
-#     # Configurar parámetros de generación
-#     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}
-
-#     # Generar respuesta de manera asíncrona
-#     def generate():
-#         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.encode('utf-8')
-
-#     return StreamingResponse(generate(), media_type="text/plain")
-
-# @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")
+    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"]})
+
+    # Preparar entrada para el modelo
+    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)
+
+    # Configurar parámetros de generación
+    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}
+
+    # Generar respuesta de manera asíncrona
+    def generate():
+        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.encode('utf-8')
+
+    return StreamingResponse(generate(), media_type="text/plain")
+
+@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")
     
-#     return simple_chat(message)
+    return simple_chat(message)
 
 with gr.Blocks(css=CSS, theme="soft", fill_height=True) as demo:
     gr.HTML(TITLE)