Spaces:
Runtime error
Runtime error
File size: 7,594 Bytes
185aa79 c5aa334 185aa79 593852b 9a3c7df 185aa79 bc504d0 185aa79 bc504d0 61931a6 c5aa334 ce8881a 61931a6 185aa79 61931a6 9b13fa2 c5aa334 ce8881a c5aa334 9b13fa2 61931a6 c5aa334 ce8881a 185aa79 097cb9f ce8881a 185aa79 ce8881a 9a3c7df bc504d0 9a3c7df bc504d0 c947682 bc504d0 c947682 bc504d0 185aa79 bc504d0 9a3c7df c947682 9a3c7df bc504d0 9a3c7df c947682 185aa79 097cb9f ce8881a 185aa79 61931a6 185aa79 9a3c7df 185aa79 61931a6 185aa79 efb1452 185aa79 5e668f8 9b13fa2 185aa79 9b13fa2 ce8881a 185aa79 151727a 185aa79 097cb9f 185aa79 61931a6 185aa79 61931a6 185aa79 58bfa53 185aa79 9b13fa2 185aa79 4680940 185aa79 151727a 185aa79 097cb9f ce8881a 185aa79 61931a6 151727a 185aa79 3e53d56 185aa79 3e53d56 185aa79 151727a 185aa79 ce8881a 185aa79 ce8881a 185aa79 ce8881a 185aa79 ce8881a 185aa79 ce8881a 185aa79 |
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 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 |
import base64
import pathlib
import gradio as gr
import spaces
import torch
from colpali_engine.models import ColPali, ColPaliProcessor
from transformers.utils.import_utils import is_flash_attn_2_available
from pdf2image import convert_from_path
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import AutoProcessor, LlavaForConditionalGeneration
PIXTAL_MODEL_ID = "mistral-community--pixtral-12b"
PIXTRAL_MODEL_SNAPSHOT = "c2756cbbb9422eba9f6c5c439a214b0392dfc998"
PIXTRAL_MODEL_PATH = (
pathlib.Path().home()
/ f".cache/huggingface/hub/models--{PIXTAL_MODEL_ID}/snapshots/{PIXTRAL_MODEL_SNAPSHOT}"
)
COLPALI_GEMMA_MODEL_ID = "vidore--colpaligemma-3b-pt-448-base"
COLPALI_GEMMA_MODEL_SNAPSHOT = "30ab955d073de4a91dc5a288e8c97226647e3e5a"
COLPALI_GEMMA_MODEL_PATH = (
pathlib.Path().home()
/ f".cache/huggingface/hub/models--{COLPALI_GEMMA_MODEL_ID}/snapshots/{COLPALI_GEMMA_MODEL_SNAPSHOT}"
)
COLPALI_MODEL_ID = "vidore--colpali-v1.3"
COLPALI_MODEL_SNAPSHOT = "1b5c8929330df1a66de441a9b5409a878f0de5b0"
COLPALI_MODEL_PATH = (
pathlib.Path().home()
/ f".cache/huggingface/hub/models--{COLPALI_MODEL_ID}/snapshots/{COLPALI_MODEL_SNAPSHOT}"
)
def image_to_base64(image_path):
with open(image_path, "rb") as img:
encoded_string = base64.b64encode(img.read()).decode("utf-8")
return f"data:image/jpeg;base64,{encoded_string}"
@spaces.GPU(duration=120)
def pixtral_inference(
images,
text,
):
if len(images) == 0:
raise gr.Error("No images for generation")
if text == "":
raise gr.Error("No query for generation")
print("LOADING MODEL")
model = LlavaForConditionalGeneration.from_pretrained(
PIXTRAL_MODEL_PATH, device_map="cuda"
)
print("LOADING MODEL DONE")
print("LOADING PROCESSOR")
processor = AutoProcessor.from_pretrained(PIXTRAL_MODEL_PATH, use_fast=True)
print("LOADING PROCESSOR DONE")
chat = [
{
"role": "user",
"content": [{"type": "image", "url": image_to_base64(i[0])} for i in images]
+ [
{"type": "text", "content": text},
],
}
]
inputs = processor.apply_chat_template(
chat,
add_generation_prompt=True,
tokenize=True,
return_dict=True,
return_tensors="pt",
).to(model.device)
print("GENERATING")
generate_ids = model.generate(**inputs, max_new_tokens=256)
print("GENERATING DONE")
print("BATCH DECODE")
output = processor.batch_decode(
generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)[0]
print("BATCH DECODE DONE")
print(output)
result = output
return result
@spaces.GPU(duration=120)
def retrieve(query: str, ds, images, k):
if len(images) == 0:
raise gr.Error("No docs/images for retrieval")
if query == "":
raise gr.Error("No query for retrieval")
model = ColPali.from_pretrained(
COLPALI_GEMMA_MODEL_PATH,
torch_dtype=torch.bfloat16,
device_map="cuda",
attn_implementation=(
"flash_attention_2" if is_flash_attn_2_available() else None
),
).eval()
model.load_adapter(COLPALI_MODEL_PATH)
model = model.eval()
processor = ColPaliProcessor.from_pretrained(COLPALI_MODEL_PATH, use_fast=True)
qs = []
with torch.no_grad():
batch_query = processor.process_queries([query]).to("cuda")
embeddings_query = model.forward(**batch_query)
qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
scores = processor.score_multi_vector(qs, ds).numpy()
top_k_indices = scores.argsort(axis=1)[0][-k:][::-1]
results = []
for idx in top_k_indices:
results.append((images[idx], f"Score {scores[0][idx]:.2f}"))
del model
del processor
torch.cuda.empty_cache()
return results
def index(files, ds):
images = convert_files(files)
return index_gpu(images, ds)
def convert_files(files):
images = []
for f in files:
images.extend(convert_from_path(f, thread_count=4))
if len(images) >= 150:
raise gr.Error("The number of images in the dataset should be less than 150.")
return images
@spaces.GPU(duration=120)
def index_gpu(images, ds):
model = ColPali.from_pretrained(
COLPALI_GEMMA_MODEL_PATH,
torch_dtype=torch.bfloat16,
device_map="cuda",
).eval()
model.load_adapter(COLPALI_MODEL_PATH)
model = model.eval()
processor = ColPaliProcessor.from_pretrained(COLPALI_MODEL_PATH, use_fast=True)
# run inference - docs
dataloader = DataLoader(
images,
batch_size=4,
shuffle=False,
collate_fn=lambda x: processor.process_images(x),
)
for batch_doc in tqdm(dataloader):
with torch.no_grad():
batch_doc = {k: v.to("cuda") for k, v in batch_doc.items()}
embeddings_doc = model(**batch_doc)
ds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
del model
del processor
torch.cuda.empty_cache()
return f"Uploaded and converted {len(images)} pages", ds, images
def get_example():
return [
[["plants_and_people.pdf"], "What is the global population in 2050 ? "],
[["plants_and_people.pdf"], "Where was Teosinte domesticated ?"],
]
css = """
#title-container {
margin: 0 auto;
max-width: 800px;
text-align: center;
}
#col-container {
margin: 0 auto;
max-width: 600px;
}
"""
file = gr.File(
file_types=[".pdf"], type="filepath", file_count="multiple", label="PDFs"
)
query = gr.Textbox("", placeholder="Enter your query here", label="Query")
with gr.Blocks(
title="Document Question Answering with ColPali & Pixtral",
theme=gr.themes.Soft(),
css=css,
) as demo:
with gr.Row(elem_id="title-container"):
gr.Markdown("""# Document Question Answering with ColPali & Pixtral""")
with gr.Column(elem_id="col-container"):
with gr.Row():
gr.Examples(
examples=get_example(),
inputs=[file, query],
)
with gr.Row():
with gr.Column(scale=2):
gr.Markdown("## Index PDFs")
file.render()
convert_button = gr.Button("π Run", variant="primary")
message = gr.Textbox("Files not yet uploaded", label="Status")
embeds = gr.State(value=[])
imgs = gr.State(value=[])
img_chunk = gr.State(value=[])
with gr.Column(scale=3):
gr.Markdown("## Retrieve with ColPali and answer with Pixtral")
query.render()
k = gr.Slider(
minimum=1,
maximum=4,
step=1,
label="Number of docs to retrieve",
value=1,
)
answer_button = gr.Button("π Run", variant="primary")
output_gallery = gr.Gallery(
label="Retrieved docs", height=400, show_label=True, interactive=False
)
output = gr.Textbox(label="Answer", lines=2, interactive=False)
convert_button.click(
index, inputs=[file, embeds], outputs=[message, embeds, imgs]
)
answer_button.click(
retrieve, inputs=[query, embeds, imgs, k], outputs=[output_gallery]
).then(pixtral_inference, inputs=[output_gallery, query], outputs=[output])
if __name__ == "__main__":
demo.queue(max_size=10).launch()
|