Update app.py

app.py

@@ -7,21 +7,23 @@ import gradio as gr
 from gradio import FileData
 import time
 import spaces
+
 ckpt = "meta-llama/Llama-3.2-11B-Vision-Instruct"
 model = MllamaForConditionalGeneration.from_pretrained(ckpt,
     torch_dtype=torch.bfloat16).to("cuda")
 processor = AutoProcessor.from_pretrained(ckpt)
 
+SYSTEM_PROMPT = """You are a Vision Language Model specialized in interpreting and extracting data from visual documents, including timesheets, invoices, charts, and other structured or semi-structured documents.
+
+Your task is to analyze the provided visual data and respond to queries with concise answers, such as single words, numbers, or short phrases.
+These documents may include tables, labels, handwritten or printed text, and graphical elements.
 
+Focus on delivering accurate, succinct answers based on the visual and contextual information provided. Avoid additional explanation unless absolutely necessary."""
 @spaces.GPU
-def bot_streaming(message, history, max_new_tokens=250):
-    
+def bot_streaming(message, history, max_new_tokens=4048):
     txt = message["text"]
-    ext_buffer = f"{txt}"
-    
-    messages= [] 
+    messages = [{"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]}]
     images = []
-    
 
     for i, msg in enumerate(history): 
         if isinstance(msg[0], tuple):
@@ -29,35 +31,30 @@ def bot_streaming(message, history, max_new_tokens=250):
             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): # text only turn
+        elif isinstance(history[i-1][0], str) and isinstance(msg[0], str):
             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): # examples
+        if isinstance(message["files"][0], str):
             image = Image.open(message["files"][0]).convert("RGB")
-        else: # regular input
+        else:
             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 == []:
         inputs = processor(text=texts, return_tensors="pt").to("cuda")
     else:
         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)
-    generated_text = ""
     
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
@@ -65,36 +62,34 @@ def bot_streaming(message, history, max_new_tokens=250):
     
     for new_text in streamer:
         buffer += new_text
-        generated_text_without_prompt = buffer
         time.sleep(0.01)
         yield buffer
 
-
-demo = gr.ChatInterface(fn=bot_streaming, title="Multimodal Llama", 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],
+demo = gr.ChatInterface(
+    fn=bot_streaming,
+    title="Multimodal Llama",
+    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=250,
-              step=10,
-              label="Maximum number of new tokens to generate",
-          )
-        ],
-      cache_examples=False,
-      description="Try Multimodal Llama by Meta with transformers in this demo. Upload an image, and start chatting about it, or simply try one of the examples below. To learn more about Llama Vision, visit [our blog post](https://huggingface.co/blog/llama32). ",
-      stop_btn="Stop Generation", 
-      fill_height=True,
-    multimodal=True)
-    
+    textbox=gr.MultimodalTextbox(),
+    additional_inputs=[
+        gr.Slider(
+            minimum=10,
+            maximum=500,
+            value=4048,
+            step=10,
+            label="Maximum number of new tokens to generate"
+        )
+    ],
+    cache_examples=False,
+    description="Try Multimodal Llama by transformers. Upload an image and start chatting about it, or try one of the examples below.",
+    stop_btn="Stop Generation",
+    fill_height=True,
+    multimodal=True
+)
+
 demo.launch(debug=True)