wealthcoders/qwen3-vl-2B

1from transformers import Qwen3VLForConditionalGeneration, AutoProcessor2from typing import Dict, List, Any3import torch4import io5from PIL import Image6import base647import time8import uuid9 10prompt = """**Task**: 11          Analyze this document image exhaustively and output in Markdown format. 12          **Rules**:  13            - Do not add any comments, provide content only;14            - Extract ALL visible text exactly as written;15            - Preserve possible additional languages;16            - Maintain line breaks, indentation, and spacing;17            - Never translate non-English text.18            - Do not add unnecessary or additional information. Do not add any links or images. Do not add Chinese symbols.19        **Important**: the output format must be Markdown (use bold text, headlines, so on)."""20 21class EndpointHandler:22    def __init__(self, path: str = "unsloth/Qwen3-VL-2B-Instruct-unsloth-bnb-4bit"):23        # Load tokenizer and model24        self.processor = AutoProcessor.from_pretrained(path)25        self.model = Qwen3VLForConditionalGeneration.from_pretrained(path, device_map="auto")26        self.model.eval()27    28    def __call__(self, data: Dict[str, Any]) -> str:29        # Prepare your messages with image and text30        inputs = data.get("inputs")31        base64image = inputs["base64"]32 33        img_bytes = base64.b64decode(base64image)34        pil_img = Image.open(io.BytesIO(img_bytes)).convert("RGB")35 36        messages = [37            {38                "role": "user",39                "content": [40                    {"type": "image", "image": pil_img},      # pass PIL image directly41                    {"type": "text", "text": prompt},42                ]43            }44        ]45 46        # Process the input and generate a response47        inputs = self.processor.apply_chat_template(48            messages,49            tokenize=True,50            add_generation_prompt=True,51            return_dict=True,52            return_tensors="pt"53        )54        inputs = inputs.to(self.model.device)55 56        generated_ids = self.model.generate(**inputs, max_new_tokens=2048)57        generated_ids_trimmed = [58            out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)59        ]60        output_text = self.processor.batch_decode(61            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False62        )63 64        response = {65            "id": f"chatcmpl-{uuid.uuid4().hex}",66            "object": "chat.completion",67            "created": int(time.time()),68            "model": "Qwen/Qwen3-VL-8B-Instruct",69            "usage": {70                # you might compute these if you can get token counts71                "prompt_tokens": None,72                "completion_tokens": None,73                "total_tokens": None74            },75            "choices": [76                {77                    "message": {78                        "role": "assistant",79                        "content": output_text[0]80                    },81                    "finish_reason": "stop",82                    "index": 083                }84            ]85        }86 87        return response