Update app.py

app.py

@@ -7,46 +7,47 @@ import gradio as gr
 from gradio import FileData
 import time
 import spaces
-
-ckpt = "Daemontatox/DocumentCogito"
+ckpt ="Daemontatox/DocumentCogito"
 model = MllamaForConditionalGeneration.from_pretrained(ckpt,
     torch_dtype=torch.bfloat16).to("cuda")
 processor = AutoProcessor.from_pretrained(ckpt)
 
-SYSTEM_PROMPT = """You are a helpful AI assistant specialized in analyzing documents, images, and visual content. 
-Your responses should be clear, accurate, and focused on the specific details present in the provided materials. 
-When analyzing documents, pay attention to key information, formatting, and context. 
-For images, consider both obvious and subtle details that might be relevant to the user's query."""
 
 @spaces.GPU()
-def bot_streaming(message, history, max_new_tokens=2048, temperature=0.7):
+def bot_streaming(message, history, max_new_tokens=2048):
+    
     txt = message["text"]
     ext_buffer = f"{txt}"
     
-    messages = [{"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]}]
+    messages= [] 
     images = []
     
-    for i, msg in enumerate(history):
+
+    for i, msg in enumerate(history): 
         if isinstance(msg[0], tuple):
             messages.append({"role": "user", "content": [{"type": "text", "text": history[i+1][0]}, {"type": "image"}]})
             messages.append({"role": "assistant", "content": [{"type": "text", "text": history[i+1][1]}]})
             images.append(Image.open(msg[0][0]).convert("RGB"))
         elif isinstance(history[i-1], tuple) and isinstance(msg[0], str):
+            # messages are already handled
             pass
-        elif isinstance(history[i-1][0], str) and isinstance(msg[0], str):
+        elif isinstance(history[i-1][0], str) and isinstance(msg[0], str): # text only turn
             messages.append({"role": "user", "content": [{"type": "text", "text": msg[0]}]})
             messages.append({"role": "assistant", "content": [{"type": "text", "text": msg[1]}]})
 
+    # add current message
     if len(message["files"]) == 1:
-        if isinstance(message["files"][0], str):
+        
+        if isinstance(message["files"][0], str): # examples
             image = Image.open(message["files"][0]).convert("RGB")
-        else:
+        else: # regular input
             image = Image.open(message["files"][0]["path"]).convert("RGB")
         images.append(image)
         messages.append({"role": "user", "content": [{"type": "text", "text": txt}, {"type": "image"}]})
     else:
         messages.append({"role": "user", "content": [{"type": "text", "text": txt}]})
 
+
     texts = processor.apply_chat_template(messages, add_generation_prompt=True)
 
     if images == []:
@@ -55,13 +56,8 @@ def bot_streaming(message, history, max_new_tokens=2048, temperature=0.7):
         inputs = processor(text=texts, images=images, return_tensors="pt").to("cuda")
     streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True)
 
-    generation_kwargs = dict(
-        inputs,
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        temperature=temperature,  # Add temperature parameter
-        do_sample=True,  # Enable sampling for temperature to take effect
-    )
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
+    generated_text = ""
     
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
@@ -73,38 +69,32 @@ def bot_streaming(message, history, max_new_tokens=2048, temperature=0.7):
         time.sleep(0.01)
         yield buffer
 
-demo = gr.ChatInterface(
-    fn=bot_streaming, 
-    title="Document Analyzer", 
-    examples=[
-        [{"text": "Which era does this piece belong to? Give details about the era.", "files":["./examples/rococo.jpg"]}, 200, 0.7],
-        [{"text": "Where do the droughts happen according to this diagram?", "files":["./examples/weather_events.png"]}, 250, 0.7],
-        [{"text": "What happens when you take out white cat from this chain?", "files":["./examples/ai2d_test.jpg"]}, 250, 0.7],
-        [{"text": "How long does it take from invoice date to due date? Be short and concise.", "files":["./examples/invoice.png"]}, 250, 0.7],
-        [{"text": "Where to find this monument? Can you give me other recommendations around the area?", "files":["./examples/wat_arun.jpg"]}, 250, 0.7],
-    ],
-    textbox=gr.MultimodalTextbox(), 
-    additional_inputs=[
-        gr.Slider(
-            minimum=10,
-            maximum=500,
-            value=2048,
-            step=10,
-            label="Maximum number of new tokens to generate",
-        ),
-        gr.Slider(  # Add temperature slider
-            minimum=0.1,
-            maximum=2.0,
-            value=0.2,
-            step=0.1,
-            label="Temperature (0.1 = focused, 2.0 = creative)",
-        )
-    ],
-    cache_examples=False,
-    description="MllM with Temperature Control",
-    stop_btn="Stop Generation", 
-    fill_height=True,
-    multimodal=True
-)
 
+demo = gr.ChatInterface(fn=bot_streaming, title="Document Analyzer", examples=[
+    [{"text": "Which era does this piece belong to? Give details about the era.", "files":["./examples/rococo.jpg"]},
+    200],
+    [{"text": "Where do the droughts happen according to this diagram?", "files":["./examples/weather_events.png"]},
+    250],
+    [{"text": "What happens when you take out white cat from this chain?", "files":["./examples/ai2d_test.jpg"]},
+    250],
+    [{"text": "How long does it take from invoice date to due date? Be short and concise.", "files":["./examples/invoice.png"]},
+    250],
+    [{"text": "Where to find this monument? Can you give me other recommendations around the area?", "files":["./examples/wat_arun.jpg"]}, 
+    250],
+    ],
+      textbox=gr.MultimodalTextbox(), 
+      additional_inputs = [gr.Slider(
+              minimum=10,
+              maximum=500,
+              value=2048,
+              step=10,
+              label="Maximum number of new tokens to generate",
+          )
+        ],
+      cache_examples=False,
+      description="MllM ",
+      stop_btn="Stop Generation", 
+      fill_height=True,
+    multimodal=True)
+    
 demo.launch(debug=True)