Views
No views yet
1from transformers import AutoModelForVision2Seq, AutoProcessor, TextStreamer
2import torch
3from PIL import Image
4
5model_id = "vanishingradient/qwen-docs-finetuned"
6
7# Load model (4-bit, fits on 16GB VRAM)
8model = AutoModelForVision2Seq.from_pretrained(
9 model_id,
10 torch_dtype=torch.float16,
11 device_map="auto",
12 trust_remote_code=True,
13 load_in_4bit=True,
14)
15
16processor = AutoProcessor.from_pretrained(
17 model_id,
18 trust_remote_code=True
19)
20
21# --------------------------------------------------
22# PLACEHOLDER: path to your local image file
23# --------------------------------------------------
24image = Image.open("/path/to/your/document_image.png")
25
26messages = [
27 {
28 "role": "user",
29 "content": [
30 {"type": "image"},
31 {"type": "text", "text": "Convert this image to markdown format."}
32 ]
33 }
34]
35
36text = processor.apply_chat_template(
37 messages,
38 tokenize=False,
39 add_generation_prompt=True
40)
41
42inputs = processor(
43 text=[text],
44 images=[image],
45 return_tensors="pt"
46).to("cuda")
47
48streamer = TextStreamer(
49 processor.tokenizer,
50 skip_prompt=True
51)
52
53_ = model.generate(
54 **inputs,
55 streamer=streamer,
56 max_new_tokens=1024,
57 temperature=0.1,
58)