fix

app.py

@@ -1,16 +1,14 @@
-# import subprocess  # 🥲 need for flash attention in QWEN model 
+import subprocess  # 🥲
 
-# subprocess.run(
-#     "pip install flash-attn --no-build-isolation",
-#     env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
-#     shell=True,
-# )
+subprocess.run(
+    "pip install flash-attn --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
 
 import spaces
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoProcessor
-# from transformers import Qwen2VLForConditionalGeneration  # Uncomment when adding QWEN back
-# from qwen_vl_utils import process_vision_info  # Uncomment when adding QWEN back
+from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
 import torch
 import os
 import json
@@ -33,15 +31,6 @@ processor = AutoProcessor.from_pretrained(
     device_map='auto'
 )
 
-# # Load Qwen model (commented out for now)
-# qwen_model = Qwen2VLForConditionalGeneration.from_pretrained(
-#     "Qwen/Qwen2-VL-7B-Instruct",
-#     torch_dtype=torch.bfloat16,
-#     attn_implementation="flash_attention_2",
-#     device_map="auto",
-# )
-# qwen_processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
-
 class GeneralRetrievalQuery(BaseModel):
     broad_topical_query: str
     broad_topical_explanation: str
@@ -91,34 +80,6 @@ Generate the queries based on this image and provide the response in the specifi
 
 prompt, pydantic_model = get_retrieval_prompt("general")
 
-# def _prep_data_for_input_qwen(image):
-#     messages = [
-#         {
-#             "role": "user",
-#             "content": [
-#                 {
-#                     "type": "image",
-#                     "image": image,
-#                 },
-#                 {"type": "text", "text": prompt},
-#             ],
-#         }
-#     ]
-#
-#     text = qwen_processor.apply_chat_template(
-#         messages, tokenize=False, add_generation_prompt=True
-#     )
-#
-#     image_inputs, video_inputs = process_vision_info(messages)
-#
-#     return qwen_processor(
-#         text=[text],
-#         images=image_inputs,
-#         videos=video_inputs,
-#         padding=True,
-#         return_tensors="pt",
-#     )
-
 def _prep_data_for_input(image):
     return processor.process(
         images=[image],
@@ -131,7 +92,7 @@ def generate_response(image):
     inputs = {k: v.to(model.device).unsqueeze(0) for k, v in inputs.items()}
     output = model.generate_from_batch(
         inputs,
-        gr.GenerationConfig(max_new_tokens=200, stop_strings="<|endoftext|>"),
+        GenerationConfig(max_new_tokens=200, stop_token="<|endoftext|>"),
         tokenizer=processor.tokenizer
     )
     generated_tokens = output[0, inputs['input_ids'].size(1):]