Update app.py

app.py

@@ -1,60 +1,32 @@
 import torch
 from PIL import Image
-import gradio as gr
-import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import os
-from threading import Thread
-
 import pymupdf
 import docx
 from pptx import Presentation
-
 from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.responses import HTMLResponse
+from typing import List, Dict
 
 app = FastAPI()
 
-@app.post("/test/")
-async def test_endpoint(message: dict):
-    if "text" not in message:
-        raise HTTPException(status_code=400, detail="Missing 'text' in request body")
-    
-    response = {"message": f"Received your message: {message['text']}"}
-    return response
-
-
+# Model and tokenizer initialization
 MODEL_LIST = ["nikravan/glm-4vq"]
-
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 MODEL_ID = MODEL_LIST[0]
 MODEL_NAME = "GLM-4vq"
 
-TITLE = "<h1>AI CHAT DOCS</h1>"
-
-DESCRIPTION = f"""
-<center>
-<p> 
-<br>
-USANDO MODELO: <a href="https://hf.co/nikravan/glm-4vq">{MODEL_NAME}</a>
-</center>"""
-
-CSS = """
-h1 {
-    text-align: center;
-    display: block;
-}
-"""
-
-
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-
-
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=True,
+    trust_remote_code=True
+)
 
 def extract_text(path):
     return open(path, 'r').read()
 
-
 def extract_pdf(path):
     doc = pymupdf.open(path)
     text = ""
@@ -62,15 +34,10 @@ def extract_pdf(path):
         text += page.get_text()
     return text
 
-
 def extract_docx(path):
     doc = docx.Document(path)
-    data = []
-    for paragraph in doc.paragraphs:
-        data.append(paragraph.text)
-    content = '\n\n'.join(data)
-    return content
-
+    data = [paragraph.text for paragraph in doc.paragraphs]
+    return '\n\n'.join(data)
 
 def extract_pptx(path):
     prs = Presentation(path)
@@ -81,49 +48,44 @@ def extract_pptx(path):
                 text += shape.text + "\n"
     return text
 
-
 def mode_load(path):
-    choice = ""
-    file_type = path.split(".")[-1]
-    print(file_type)
-    if file_type in ["pdf", "txt", "py", "docx", "pptx", "json", "cpp", "md"]:
-        if file_type.endswith("pdf"):
+    file_type = path.split(".")[-1].lower()
+    if file_type in ["pdf", "txt", "py", "docx", "pptx"]:
+        if file_type == "pdf":
             content = extract_pdf(path)
-        elif file_type.endswith("docx"):
+        elif file_type == "docx":
             content = extract_docx(path)
-        elif file_type.endswith("pptx"):
+        elif file_type == "pptx":
             content = extract_pptx(path)
         else:
             content = extract_text(path)
-        choice = "doc"
-        print(content[:100])
-        return choice, content[:5000]
-
-
+        return "doc", content[:5000]
     elif file_type in ["png", "jpg", "jpeg", "bmp", "tiff", "webp"]:
         content = Image.open(path).convert('RGB')
-        choice = "image"
-        return choice, content
-
+        return "image", content
     else:
-        raise gr.Error("Oops, unsupported files.")
-
+        raise HTTPException(status_code=400, detail="Unsupported file type")
 
-@spaces.GPU()
-def stream_chat(message, history: list, temperature: float, max_length: int, top_p: float, top_k: int, penalty: float):
+@app.post("/test/")
+async def test_endpoint(message: Dict[str, str]):
+    if "text" not in message:
+        raise HTTPException(status_code=400, detail="Missing 'text' in request body")
     
-    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}')
+    response = {"message": f"Received your message: {message['text']}"}
+    return response
+
+@app.post("/chat/")
+async def chat_endpoint(
+    message: Dict[str, str],
+    history: List[Dict[str, str]] = [],
+    temperature: float = 0.8,
+    max_length: int = 4096,
+    top_p: float = 1.0,
+    top_k: int = 10,
+    penalty: float = 1.0
+):
     conversation = []
-    prompt_files = []
-    if message["files"]:
+    if "files" in message and message["files"]:
         choice, contents = mode_load(message["files"][-1])
         if choice == "image":
             conversation.append({"role": "user", "image": contents, "content": message['text']})
@@ -132,35 +94,26 @@ def stream_chat(message, history: list, temperature: float, max_length: int, top
             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}")
+            if len(history) > 0:
+                choice, contents = mode_load(history[-1][0])
+                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)
+    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,
@@ -168,97 +121,11 @@ def stream_chat(message, history: list, temperature: float, max_length: int, top
         top_p=top_p,
         top_k=top_k,
         temperature=temperature,
-        repetition_penalty=penalty,
-        eos_token_id=[151329, 151336, 151338],
+        repetition_penalty=penalty
     )
-    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
-
-
-chatbot = gr.Chatbot(
-    #rtl=True,
-)
-chat_input = gr.MultimodalTextbox(
-    interactive=True,
-    placeholder="Enter message or upload a file ...",
-    show_label=False,
-    #rtl=True,
-
-
-
-)
-
-EXAMPLES = [
-    [{"text": "Resumir Documento"}],
-    [{"text": "Explicar la Imagen"}],
-    [{"text": "¿De qué es la foto?", "files": ["perro.jpg"]}],
-    [{"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.", }]
-]
-
-with gr.Blocks(css=CSS, theme="soft", fill_height=True) as demo:
-    gr.HTML(TITLE)
-    gr.HTML(DESCRIPTION)
-    gr.ChatInterface(
-        fn=stream_chat,
-        multimodal=True,
-
-
-        textbox=chat_input,
-        chatbot=chatbot,
-        fill_height=True,
-        additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
-        additional_inputs=[
-            gr.Slider(
-                minimum=0,
-                maximum=1,
-                step=0.1,
-                value=0.8,
-                label="Temperature",
-                render=False,
-            ),
-            gr.Slider(
-                minimum=1024,
-                maximum=8192,
-                step=1,
-                value=4096,
-                label="Max Length",
-                render=False,
-            ),
-            gr.Slider(
-                minimum=0.0,
-                maximum=1.0,
-                step=0.1,
-                value=1.0,
-                label="top_p",
-                render=False,
-            ),
-            gr.Slider(
-                minimum=1,
-                maximum=20,
-                step=1,
-                value=10,
-                label="top_k",
-                render=False,
-            ),
-            gr.Slider(
-                minimum=0.0,
-                maximum=2.0,
-                step=0.1,
-                value=1.0,
-                label="Repetition penalty",
-                render=False,
-            ),
-        ],
-    ),
-    gr.Examples(EXAMPLES, [chat_input])
-
-if __name__ == "__main__":
-    
-    demo.queue(api_open=False).launch(show_api=False, share=False, )#server_name="0.0.0.0", )
+        return {"response": buffer}