Update app.py

app.py

@@ -226,16 +226,13 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
 
         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":
@@ -244,28 +241,12 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
                 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)
 
         streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
 
-
-        print("--------------")
-        print("  ")
-        print(input_ids)
-        print("--------------")
-        print(streamer)
-        print("  ")
-        print("--------------")
-
         generate_kwargs = dict(
             max_length=max_length,
             do_sample=True,
@@ -278,32 +259,46 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
 
         gen_kwargs = {**input_ids, **generate_kwargs}
 
-
         with torch.no_grad():
             thread = Thread(target=model.generate, kwargs=gen_kwargs)
             thread.start()
-            print("--------------")
-            print("  ")
-            print(thread)
-            print("  ")
-            print("--------------")
             buffer = ""
             for new_text in streamer:
                 buffer += new_text
-                yield new_text
-            print("--------------")
-            print(new_text)
-            print("Buffer: ")
-            print("  ")
-            print(buffer)
-            print("  ")
-            print("--------------")
 
-
-        return StreamingResponse(new_text, media_type="text/plain")
+        return buffer
 
     except Exception as e:
-        return PlainTextResponse(f"Error: {str(e)}")
+        return f"Error: {str(e)}"
+
+@app.post("/chat/")
+async def test_endpoint(
+    text: str = Form(...), 
+    file: UploadFile = File(None)
+):
+    if file:
+        file_content = BytesIO(await file.read())
+        file_name = file.filename
+        
+        message = {
+            "text": text,
+            "file_content": file_content,
+            "file_name": file_name
+        }
+    else:
+        message = {
+            "text": text,
+            "file_content": None,
+            "file_name": None
+        }
+    
+    print(message)
+    response = simple_chat(message)
+    
+    if isinstance(response, str) and response.startswith("Error:"):
+        return PlainTextResponse(response, status_code=500)
+    else:
+        return StreamingResponse(BytesIO(response.encode()), media_type="text/plain")
         
 # 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:
@@ -368,36 +363,36 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
 
 
 
-@app.post("/chat/")
-async def test_endpoint(
-    text: str = Form(...), 
-    file: UploadFile = File(None)
-):
-    # Verificar si se ha subido un archivo
-    if file:
-        # Leer el archivo en memoria
-        file_content = BytesIO(await file.read())
-        file_name = file.filename
+# @app.post("/chat/")
+# async def test_endpoint(
+#     text: str = Form(...), 
+#     file: UploadFile = File(None)
+# ):
+#     # Verificar si se ha subido un archivo
+#     if file:
+#         # Leer el archivo en memoria
+#         file_content = BytesIO(await file.read())
+#         file_name = file.filename
         
-        # Construir el mensaje con el archivo y el texto
-        message = {
-            "text": text,
-            "file_content": file_content,
-            "file_name": file_name
-        }
-    else:
-        # Si no se sube archivo, solo se incluye el texto
-        message = {
-            "text": text,
-            "file_content": None,
-            "file_name": None
-        }
+#         # Construir el mensaje con el archivo y el texto
+#         message = {
+#             "text": text,
+#             "file_content": file_content,
+#             "file_name": file_name
+#         }
+#     else:
+#         # Si no se sube archivo, solo se incluye el texto
+#         message = {
+#             "text": text,
+#             "file_content": None,
+#             "file_name": None
+#         }
     
-    # Llamar a la función `simple_chat` con el mensaje
-    print(message)
-    response = simple_chat(message)
+#     # Llamar a la función `simple_chat` con el mensaje
+#     print(message)
+#     response = simple_chat(message)
     
-    return response
+#     return response
 
 
 with gr.Blocks(css=CSS, theme="soft", fill_height=True) as demo: