Update app.py

app.py

@@ -1,158 +1,100 @@
-
-import subprocess 
-
-subprocess.run(
-    'pip install --force-reinstall --no-cache-dir --no-deps git+https://github.com/unslothai/unsloth.git ',
-    shell=True
-)
-
-
-subprocess.run(
-    'pip install unsloth_zoo',
-    shell=True
-)
-
-from transformers import AutoTokenizer, TextStreamer
+from transformers import MllamaForConditionalGeneration, AutoProcessor, TextIteratorStreamer
 from PIL import Image
+import requests
 import torch
 from threading import Thread
 import gradio as gr
 from gradio import FileData
 import time
 import spaces
-from unsloth import FastVisionModel
-
-
-# Load model and tokenizer
 ckpt = "Daemontatox/DocumentLlama"
-model, tokenizer = FastVisionModel.from_pretrained(
-    ckpt,
-    load_in_4bit=True,
-    use_gradient_checkpointing="unsloth",
-)
+model = MllamaForConditionalGeneration.from_pretrained(ckpt,
+    torch_dtype=torch.bfloat16,load_in_4bit=True).to("cuda")
+processor = AutoProcessor.from_pretrained(ckpt)
 
-# Enable inference mode
-FastVisionModel.for_inference(model)
 
 @spaces.GPU()
 def bot_streaming(message, history, max_new_tokens=2048):
+    
     txt = message["text"]
-    messages = []
+    ext_buffer = f"{txt}"
+    
+    messages= [] 
     images = []
     
-    # Process history
-    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]}]
-            })
+            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):
-            messages.append({
-                "role": "user",
-                "content": [{"type": "text", "text": msg[0]}]
-            })
-            messages.append({
-                "role": "assistant",
-                "content": [{"type": "text", "text": msg[1]}]
-            })
+        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]}]})
 
-    # Handle current message
+    # add current message
     if len(message["files"]) == 1:
-        if isinstance(message["files"][0], str):  # examples
+        
+        if isinstance(message["files"][0], str): # examples
             image = Image.open(message["files"][0]).convert("RGB")
-        else:  # regular input
+        else: # regular input
             image = Image.open(message["files"][0]["path"]).convert("RGB")
         images.append(image)
-        messages.append({
-            "role": "user",
-            "content": [
-                {"type": "image"},
-                {"type": "text", "text": txt}
-            ]
-        })
+        messages.append({"role": "user", "content": [{"type": "text", "text": txt}, {"type": "image"}]})
     else:
-        messages.append({
-            "role": "user",
-            "content": [{"type": "text", "text": txt}]
-        })
+        messages.append({"role": "user", "content": [{"type": "text", "text": txt}]})
 
-    # Prepare inputs
-    input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
-    
-    if images:
-        inputs = tokenizer(
-            images[-1],  # Use the last image
-            input_text,
-            add_special_tokens=False,
-            return_tensors="pt"
-        ).to("cuda")
-    else:
-        inputs = tokenizer(
-            input_text,
-            add_special_tokens=False,
-            return_tensors="pt"
-        ).to("cuda")
 
-    # Setup streaming
-    text_streamer = TextStreamer(tokenizer, skip_prompt=True)
-    buffer = ""
+    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)
 
-    def generate():
-        nonlocal buffer
-        output_ids = model.generate(
-            **inputs,
-            streamer=text_streamer,
-            max_new_tokens=max_new_tokens,
-            use_cache=True,
-            temperature=1.5,
-            min_p=0.1
-        )
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
+    generated_text = ""
     
-    thread = Thread(target=generate)
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
+    buffer = ""
     
-    for new_text in text_streamer:
+    for new_text in streamer:
         buffer += new_text
+        generated_text_without_prompt = buffer
         time.sleep(0.01)
         yield buffer
 
-# Setup Gradio interface
-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 = 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)