feazrgf

app.py

@@ -13,7 +13,83 @@ from pdf2image import convert_from_path
 from PIL import Image
 from torch.utils.data import DataLoader
 from tqdm import tqdm
-from transformers import AutoProcessor
+from transformers import AutoProcessor, Idefics3ForConditionalGeneration
+import re
+import time
+from PIL import Image
+import torch
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+
+## Load idefics
+id_processor = AutoProcessor.from_pretrained("HuggingFaceM4/Idefics3-8B-Llama3")
+
+id_model = Idefics3ForConditionalGeneration.from_pretrained("HuggingFaceM4/Idefics3-8B-Llama3", 
+        torch_dtype=torch.bfloat16,
+        #_attn_implementation="flash_attention_2"
+                                                        ).to("cuda")
+
+BAD_WORDS_IDS = processor.tokenizer(["<image>", "<fake_token_around_image>"], add_special_tokens=False).input_ids
+EOS_WORDS_IDS = [processor.tokenizer.eos_token_id]
+
+@spaces.GPU
+def model_inference(
+    images, text, assistant_prefix= None, decoding_strategy = "Greedy", temperature= 0.4, max_new_tokens=512,
+    repetition_penalty=1.2, top_p=0.8
+):
+    if text == "" and not images:
+        gr.Error("Please input a query and optionally image(s).")
+
+    if text == "" and images:
+        gr.Error("Please input a text query along the image(s).")
+
+    if isinstance(images, Image.Image):
+        images = [images]
+
+
+    resulting_messages = [
+                {
+                    "role": "user",
+                    "content": [{"type": "image"}] + [
+                        {"type": "text", "text": text}
+                    ]
+                }
+            ]
+
+    if assistant_prefix:
+      text = f"{assistant_prefix} {text}"
+
+
+    prompt = processor.apply_chat_template(resulting_messages, add_generation_prompt=True)
+    inputs = processor(text=prompt, images=[images], return_tensors="pt")
+    inputs = {k: v.to("cuda") for k, v in inputs.items()}
+
+    generation_args = {
+        "max_new_tokens": max_new_tokens,
+        "repetition_penalty": repetition_penalty,
+
+    }
+
+    assert decoding_strategy in [
+        "Greedy",
+        "Top P Sampling",
+    ]
+    if decoding_strategy == "Greedy":
+        generation_args["do_sample"] = False
+    elif decoding_strategy == "Top P Sampling":
+        generation_args["temperature"] = temperature
+        generation_args["do_sample"] = True
+        generation_args["top_p"] = top_p
+
+
+    generation_args.update(inputs)
+
+    # Generate
+    generated_ids = model.generate(**generation_args)
+
+    generated_texts = processor.batch_decode(generated_ids[:, generation_args["input_ids"].size(1):], skip_special_tokens=True)
+    return generated_texts[0]
 
 # Load model
 model_name = "vidore/colpali-v1.2"
@@ -96,7 +172,7 @@ def index_gpu(images, ds):
             embeddings_doc = model(**batch_doc)
         ds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
     return f"Uploaded and converted {len(images)} pages", ds, images
-    
+
 @spaces.GPU
 def answer_gpu():
     return 0
@@ -116,6 +192,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             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("## 2️⃣ Search")
@@ -133,10 +210,11 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     output_gallery = gr.Gallery(label="Retrieved Documents", height=600, show_label=True)
 
     convert_button.click(index, inputs=[file, embeds], outputs=[message, embeds, imgs])
-    search_button.click(search, inputs=[query, embeds, imgs, k], outputs=[output_gallery])
+    search_button.click(search, inputs=[query, embeds, imgs, k], outputs=[output_gallery, img_chunk])
 
     answer_button = gr.Button("Answer", variant="primary")
-    answer_button.click(answer_gpu, inputs=[])
+    output = gr.Textbox(label="Output")
+    answer_button.click(model_inference, inputs=[img_chunk, query], outputs=output)
 
 if __name__ == "__main__":
     demo.queue(max_size=10).launch(debug=True)