Initial commit

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,20 @@
+---
+title: ColQwen & Pixtral
+short_description: Document Question Answering with ColQwen & Pixtral
+emoji: 👀
+colorFrom: purple
+colorTo: blue
+sdk: gradio
+sdk_version: 4.44.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+models:
+  - vidore--colqwen2-base
+  - vidore/colqwen2-v0.1
+  - mistral-community/pixtral-12b-240910
+preload_from_hub:
+  - vidore/colqwen2-base added_tokens.json,chat_template.json,config.json,generation_config.json,merges.txt,model-00001-of-00002.safetensors,model-00002-of-00002.safetensors,model.safetensors.index.json,preprocessor_config.json,special_tokens_map.json,tokenizer.json,tokenizer_config.json,vocab.json c722b912b50b14e404b91679db710fa2e1c6a762
+  - vidore/colqwen2-v0.1 adapter_config.json,adapter_model.safetensors,added_tokens.json,chat_template.json,preprocessor_config.json,special_tokens_map.json,tokenizer.json,tokenizer_config.json,vocab.json 6b9ef3c32c97c0bb3be99bc35a05d9f30e0cada5
+  - mistral-community/pixtral-12b-240910 params.json,tekken.json,consolidated.safetensors 95758896fcf4691ec9674f29ec90d1441d9d26d2
+---

app.py

@@ -0,0 +1,242 @@
+import base64
+from typing import cast
+import pathlib
+import gradio as gr
+import spaces
+import torch
+from ColQwen_engine.models import ColQwen2, ColQwen2Processor
+from mistral_common.protocol.instruct.messages import (
+    ImageURLChunk,
+    TextChunk,
+    UserMessage,
+)
+from mistral_common.protocol.instruct.request import ChatCompletionRequest
+from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+from mistral_inference.generate import generate
+from mistral_inference.transformer import Transformer
+from pdf2image import convert_from_path
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+PIXTAL_MODEL_ID = "mistral-community--pixtral-12b-240910"
+PIXTRAL_MODEL_SNAPSHOT = "95758896fcf4691ec9674f29ec90d1441d9d26d2"
+PIXTRAL_MODEL_PATH = (
+    pathlib.Path().home()
+    / f".cache/huggingface/hub/models--{PIXTAL_MODEL_ID}/snapshots/{PIXTRAL_MODEL_SNAPSHOT}"
+)
+
+COLQWEN_BASE_MODEL_ID = "vidore--colqwen2-base"
+COLQWEN_BASE_MODEL_SNAPSHOT = "c722b912b50b14e404b91679db710fa2e1c6a762"
+COLQWEN_BASE_MODEL_PATH = (
+    pathlib.Path().home()
+    / f".cache/huggingface/hub/models--{COLQWEN_BASE_MODEL_ID}/snapshots/{COLQWEN_BASE_MODEL_SNAPSHOT}"
+)
+COLQWEN_MODEL_ID = "vidore--colqwen2-v0.1"
+COLQWEN_MODEL_SNAPSHOT = "6b9ef3c32c97c0bb3be99bc35a05d9f30e0cada5"
+COLQWEN_MODEL_PATH = (
+    pathlib.Path().home()
+    / f".cache/huggingface/hub/models--{COLQWEN_MODEL_ID}/snapshots/{COLQWEN_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=60)
+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")
+    tokenizer = MistralTokenizer.from_file(f"{PIXTRAL_MODEL_PATH}/tekken.json")
+    model = Transformer.from_folder(PIXTRAL_MODEL_PATH)
+
+    messages = [
+        UserMessage(
+            content=[ImageURLChunk(image_url=image_to_base64(i[0])) for i in images]
+            + [TextChunk(text=text)]
+        )
+    ]
+
+    completion_request = ChatCompletionRequest(messages=messages)
+
+    encoded = tokenizer.encode_chat_completion(completion_request)
+
+    images = encoded.images
+    tokens = encoded.tokens
+
+    out_tokens, _ = generate(
+        [tokens],
+        model,
+        images=[images],
+        max_tokens=512,
+        temperature=0.45,
+        eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id,
+    )
+    result = tokenizer.decode(out_tokens[0])
+    return result
+
+
+@spaces.GPU(duration=60)
+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 = ColQwen2.from_pretrained(
+        COLQWEN_BASE_MODEL_PATH,
+        torch_dtype=torch.bfloat16,
+        device_map="cuda",
+    ).eval()
+
+    model.load_adapter(COLQWEN_MODEL_PATH)
+    model = model.eval()
+    processor = cast(
+        ColQwen2Processor, ColQwen2Processor.from_pretrained(COLQWEN_MODEL_PATH)
+    )
+
+    qs = []
+    with torch.no_grad():
+        batch_query = processor.process_queries([query])
+        batch_query = {k: v.to("cuda") for k, v in batch_query.items()}
+        embeddings_query = model(**batch_query)
+        qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
+
+    scores = processor.score(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=60)
+def index_gpu(images, ds):
+    model = ColQwen2.from_pretrained(
+        COLQWEN_BASE_MODEL_PATH,
+        torch_dtype=torch.bfloat16,
+        device_map="cuda",
+    ).eval()
+
+    model.load_adapter(COLQWEN_MODEL_PATH)
+    model = model.eval()
+    processor = cast(
+        ColQwen2Processor, ColQwen2Processor.from_pretrained(COLQWEN_MODEL_PATH)
+    )
+
+    # 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"], file_count="multiple", label="PDFs")
+query = gr.Textbox("", placeholder="Enter your query here", label="Query")
+
+with gr.Blocks(
+    title="Document Question Answering with ColQwen & Pixtral",
+    theme=gr.themes.Soft(),
+    css=css,
+) as demo:
+    with gr.Row(elem_id="title-container"):
+        gr.Markdown("""# Document Question Answering with ColQwen & 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 ColQwen 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()

packages.txt

@@ -0,0 +1 @@
+poppler-utils

requirements.txt

@@ -0,0 +1,7 @@
+gradio==4.44.0
+transformers==4.45.1
+huggingface_hub==0.25.0
+pdf2image==1.17.0
+spaces==0.30.2
+colpali_engine==0.3.1
+mistral_inference==1.4.0