Jwalit/gemma4-e4b-kyc-document-extractor

import torch
from transformers import AutoProcessor, AutoModelForImageTextToText
from PIL import Image

model_id = "Jwalit/gemma4-e4b-kyc-document-extractor"
processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForImageTextToText.from_pretrained(
    model_id, device_map="auto", torch_dtype=torch.bfloat16
)

image = Image.open("document.jpg").convert("RGB")

messages = [
    {"role": "system", "content": [{"type": "text", "text": "You are an expert KYC document analyst. Always respond with accurate, structured JSON output."}]},
    {"role": "user", "content": [
        {"type": "image"},
        {"type": "text", "text": "Classify this document and extract all information as structured JSON."}
    ]}
]

inputs = processor.apply_chat_template(
    messages, add_generation_prompt=True, tokenize=True,
    return_dict=True, return_tensors="pt", images=[image]
).to(model.device)

with torch.no_grad():
    output = model.generate(**inputs, max_new_tokens=1024, temperature=0.1)
    
result = processor.batch_decode(output[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True)[0]
print(result)

# Start OpenAI-compatible server
python -m vllm.entrypoints.openai.api_server \
    --model Jwalit/gemma4-e4b-kyc-document-extractor \
    --trust-remote-code \
    --max-model-len 4096 \
    --dtype bfloat16 \
    --gpu-memory-utilization 0.9

from openai import OpenAI
import base64

client = OpenAI(base_url="http://localhost:8000/v1", api_key="dummy")

with open("document.jpg", "rb") as f:
    img_b64 = base64.b64encode(f.read()).decode()

response = client.chat.completions.create(
    model="Jwalit/gemma4-e4b-kyc-document-extractor",
    messages=[
        {"role": "system", "content": "You are an expert KYC document analyst. Always respond with accurate, structured JSON output."},
        {"role": "user", "content": [
            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}},
            {"type": "text", "text": "Classify and extract all fields from this KYC document as JSON."}
        ]}
    ],
    max_tokens=1024,
    temperature=0.1
)
print(response.choices[0].message.content)

from vllm import LLM, SamplingParams

llm = LLM(
    model="Jwalit/gemma4-e4b-kyc-document-extractor",
    trust_remote_code=True,
    max_model_len=4096,
    dtype="bfloat16",
)

sampling_params = SamplingParams(temperature=0.1, max_tokens=1024)
# Use llm.chat() with image messages for batch processing

Gemma4ForConditionalGeneration
├── Vision Encoder (SigLIP, FROZEN)
│   ├── 16 layers, 768-dim, 12 attention heads
│   ├── Patch size: 16, Pooling kernel: 3
│   └── Output: 280 soft tokens per image
├── Text Decoder (LoRA applied here)
│   ├── 42 layers (36 sliding + 6 full attention)
│   ├── 2560 hidden, 8 heads, GQA
│   ├── 262K vocab, 131K context
│   └── LoRA on: q/k/v/o_proj + gate/up/down_proj
└── Audio Encoder (unused, frozen)

# Install dependencies
pip install torch transformers trl datasets peft accelerate bitsandbytes trackio flash-attn pillow

# Run training (requires GPU with ≥24GB VRAM, recommended: A100 80GB)
python train_kyc_vlm.py

trl sft \
    --model_name_or_path google/gemma-4-E4B-it \
    --dataset_name Jwalit/kyc-document-extraction-vlm \
    --output_dir ./gemma4-kyc-extractor \
    --learning_rate 2e-4 \
    --num_train_epochs 3 \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --bf16 \
    --gradient_checkpointing \
    --push_to_hub \
    --hub_model_id Jwalit/gemma4-e4b-kyc-document-extractor

from peft import AutoPeftModelForCausalLM
import torch

model = AutoPeftModelForCausalLM.from_pretrained(
    "Jwalit/gemma4-e4b-kyc-document-extractor",
    device_map="auto",
    torch_dtype=torch.bfloat16,
)
merged_model = model.merge_and_unload()
merged_model.save_pretrained("./merged-kyc-extractor")
# Then push merged model for faster vLLM serving

{
  "document_type": "aadhaar_card",
  "full_name": "Rajesh Kumar Singh",
  "date_of_birth": "15/03/1985",
  "gender": "Male",
  "father_name": "Suresh Kumar Singh",
  "aadhaar_number": "1234 5678 9012",
  "address": "123, MG Road, Mumbai, Maharashtra - 400001",
  "vid": "1234 5678 9012 3456"
}