Spaces:
Running
Running
File size: 4,462 Bytes
15245b5 686ef17 15d82cf 15245b5 15d82cf 15245b5 15d82cf 15245b5 686ef17 121a196 0318ee3 686ef17 15245b5 686ef17 15245b5 686ef17 15245b5 686ef17 15245b5 686ef17 15245b5 686ef17 15245b5 686ef17 15245b5 686ef17 15245b5 686ef17 15245b5 686ef17 15245b5 686ef17 15245b5 686ef17 15245b5 686ef17 15245b5 686ef17 15245b5 686ef17 15245b5 686ef17 15245b5 686ef17 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
from transformers import AutoTokenizer, TextStreamer
from PIL import Image
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",
)
# Enable inference mode
FastVisionModel.for_inference(model)
@spaces.GPU()
def bot_streaming(message, history, max_new_tokens=2048):
txt = message["text"]
messages = []
images = []
# Process 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):
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]}]
})
# Handle current message
if len(message["files"]) == 1:
if isinstance(message["files"][0], str): # examples
image = Image.open(message["files"][0]).convert("RGB")
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}
]
})
else:
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 = ""
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
)
thread = Thread(target=generate)
thread.start()
for new_text in text_streamer:
buffer += new_text
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.launch(debug=True) |