Update app.py

app.py

@@ -11,7 +11,7 @@ import pymupdf
 import docx
 from pptx import Presentation
 
-from fastapi import FastAPI, File, UploadFile, HTTPException, Form
+from fastapi import FastAPI, File, Form, UploadFile, HTTPException
 from fastapi.responses import HTMLResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import StreamingResponse
@@ -210,7 +210,61 @@ EXAMPLES = [
 
 
 # Definir la función simple_chat
-@spaces.GPU()
+# @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
+
+
 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(
@@ -223,12 +277,14 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
         conversation = []
 
         if "file" in message and message["file"]:
-            file_path = message["file"]
-            choice, contents = mode_load(file_path)
+            file = message["file"]
+            # Leer el archivo recibido
+            file_contents = io.BytesIO(file).read()
+            choice, contents = mode_load(file_contents)
             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"]
+                format_msg = contents + "\n\n\n" + "{} files uploaded.\n".format(file.filename) + message["text"]
                 conversation.append({"role": "user", "content": format_msg})
         else:
             conversation.append({"role": "user", "content": message["text"]})
@@ -253,16 +309,25 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
     except Exception as e:
         return PlainTextResponse(f"Error: {str(e)}")
 
+
 @app.post("/chat/")
-async def test_endpoint(message: dict):
-    if "text" not in message:
+async def test_endpoint(
+    text: str = Form(...),
+    file: UploadFile = File(None)  # 'None' makes the file optional
+):
+    if not text:
         raise HTTPException(status_code=400, detail="Missing 'text' in request body")
 
-    if "file" not in message:
-        print("Sin File")
+    if file:
+        # Process the file if it is provided
+        file_contents = await file.read()
+        # Do something with file_contents if needed
+        print("File received")
+    else:
+        print("No file received")
     
-    response = simple_chat(message)
-    return response
+    response = simple_chat(text)
+    return {"response": response}
 
 
 with gr.Blocks(css=CSS, theme="soft", fill_height=True) as demo: