wehub-resource-sync
69 lines
2.0 KiB
Python
69 lines
2.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Model loading: Qwen3-VL + LoRA for embedding fine-tuning."""
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from peft import LoraConfig, get_peft_model
|
|
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
|
|
|
|
|
|
def load_model_for_training(
|
|
model_path: str,
|
|
gpu_id: int,
|
|
lora_r: int = 32,
|
|
lora_alpha: int = 32,
|
|
lora_dropout: float = 0.05,
|
|
):
|
|
"""Load Qwen3-VL with LoRA adapters for fine-tuning.
|
|
|
|
Uses Qwen3VLForConditionalGeneration (NOT AutoModel, which loads
|
|
Qwen3VLModel with random language_model weights).
|
|
"""
|
|
model = Qwen3VLForConditionalGeneration.from_pretrained(
|
|
model_path,
|
|
dtype=torch.bfloat16,
|
|
trust_remote_code=True,
|
|
)
|
|
lora_config = LoraConfig(
|
|
r=lora_r,
|
|
lora_alpha=lora_alpha,
|
|
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
|
|
lora_dropout=lora_dropout,
|
|
task_type="FEATURE_EXTRACTION",
|
|
)
|
|
model = get_peft_model(model, lora_config)
|
|
model.print_trainable_parameters()
|
|
device = f"cuda:{gpu_id}"
|
|
model = model.to(device)
|
|
return model
|
|
|
|
|
|
def load_processor(
|
|
model_path: str, min_pixels: int = 128 * 28 * 28, max_pixels: int = 256 * 28 * 28
|
|
):
|
|
"""Load the processor with configurable image resolution.
|
|
|
|
Default resolution is reduced from the model default (~1.3M pixels) to
|
|
~100K-200K pixels for training speed. Full resolution can be restored
|
|
for inference by setting max_pixels=1280*28*28.
|
|
"""
|
|
return AutoProcessor.from_pretrained(
|
|
model_path,
|
|
trust_remote_code=True,
|
|
min_pixels=min_pixels,
|
|
max_pixels=max_pixels,
|
|
)
|
|
|
|
|
|
def pool_and_normalize(hidden_states, attention_mask):
|
|
"""Last-token pooling + L2 normalization.
|
|
|
|
Matches the production embedding pipeline (embed_tiles.py direct_gpu backend).
|
|
"""
|
|
last_idx = attention_mask.sum(dim=1) - 1
|
|
pooled = hidden_states[
|
|
torch.arange(hidden_states.size(0), device=hidden_states.device),
|
|
last_idx,
|
|
]
|
|
return F.normalize(pooled, p=2, dim=-1)
|