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()