Use pixtral with hf

README.md

@@ -16,5 +16,5 @@ models:
 preload_from_hub:
   - vidore/colpaligemma-3b-pt-448-base config.json,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 30ab955d073de4a91dc5a288e8c97226647e3e5a
   - vidore/colpali-v1.3 adapter_config.json,adapter_model.safetensors,preprocessor_config.json,special_tokens_map.json,tokenizer.json,tokenizer_config.json 1b5c8929330df1a66de441a9b5409a878f0de5b0
-  - mistral-community/pixtral-12b-240910 params.json,tekken.json,consolidated.safetensors 59794e97cb4f322f6223bb0d57b4d7523f0e27c6
+  - mistral-community/pixtral-12b chat_template.json,config.json,generation_config.json,model-00001-of-00006.safetensors,model-00002-of-00006.safetensors,model-00003-of-00006.safetensors,model-00004-of-00006.safetensors,model-00005-of-00006.safetensors,model-00006-of-00006.safetensors,model.safetensors.index.json,preprocessor_config.json,processor_config.json,special_tokens_map.json,tokenizer.json,tokenizer_config.json c2756cbbb9422eba9f6c5c439a214b0392dfc998
 ---

app.py

@@ -4,21 +4,14 @@ import gradio as gr
 import spaces
 import torch
 from colpali_engine.models import ColPali, ColPaliProcessor
-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
+from transformers import AutoProcessor, LlavaForConditionalGeneration
+
 
-PIXTAL_MODEL_ID = "mistral-community--pixtral-12b-240910"
-PIXTRAL_MODEL_SNAPSHOT = "59794e97cb4f322f6223bb0d57b4d7523f0e27c6"
+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}"
@@ -54,32 +47,33 @@ def pixtral_inference(
         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, dtype=torch.bfloat16)
-
-    messages = [
-        UserMessage(
-            content=[ImageURLChunk(image_url=image_to_base64(i[0])) for i in images]
-            + [TextChunk(text=text)]
-        )
+    
+    model = LlavaForConditionalGeneration.from_pretrained(PIXTRAL_MODEL_PATH)
+    processor = AutoProcessor.from_pretrained(PIXTRAL_MODEL_PATH, use_fast=True)
+
+    chat = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "content": text},
+            ]
+            + [{"type": "image", "url": image_to_base64(i[0])} for i in images],
+        }
     ]
 
-    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])
+    inputs = processor.apply_chat_template(
+        chat,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt",
+    ).to(model.device)
+    generate_ids = model.generate(**inputs, max_new_tokens=500)
+    output = processor.batch_decode(
+        generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )[0]
+
+    result = output.text
     return result
 
 

requirements.txt

@@ -3,6 +3,4 @@ transformers==4.53.2
 huggingface_hub==0.33.4
 pdf2image==1.17.0
 spaces==0.37.1
-colpali_engine==0.3.11
-mistral_inference==1.6.0
-mistral_common[opencv]==1.7.0
+colpali_engine==0.3.11