314 lines
12 KiB
Python
314 lines
12 KiB
Python
|
import base64
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import math
|
||
|
|
from io import BytesIO
|
||
|
|
from typing import Any, Dict, List, Literal, Optional, Union
|
||
|
|
from urllib.parse import urlparse
|
||
|
|
|
||
|
|
import requests
|
||
|
|
import torch
|
||
|
|
from PIL import Image
|
||
|
|
from torch import nn
|
||
|
|
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
|
||
|
|
|
||
|
|
class Transformer(nn.Module):
|
||
|
|
save_in_root: bool = True
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
model_name_or_path: str = 'llamaindex/vdr-2b-multi-v1',
|
||
|
|
processor_name_or_path: Optional[str] = None,
|
||
|
|
max_pixels: int = 768 * 28 * 28,
|
||
|
|
min_pixels: int = 1 * 28 * 28,
|
||
|
|
dimension: int = 2048,
|
||
|
|
max_seq_length: Optional[int] = None,
|
||
|
|
model_args: Optional[Dict[str, Any]] = None,
|
||
|
|
processor_args: Optional[Dict[str, Any]] = None,
|
||
|
|
tokenizer_args: Optional[Dict[str, Any]] = None,
|
||
|
|
config_args: Optional[Dict[str, Any]] = None,
|
||
|
|
cache_dir: Optional[str] = None,
|
||
|
|
backend: Literal['torch', 'onnx', 'openvino'] = 'torch',
|
||
|
|
**kwargs,
|
||
|
|
) -> None:
|
||
|
|
super(Transformer, self).__init__()
|
||
|
|
|
||
|
|
if backend != 'torch':
|
||
|
|
raise ValueError(
|
||
|
|
f'Backend \'{backend}\' is not supported, please use \'torch\' instead'
|
||
|
|
)
|
||
|
|
|
||
|
|
self.dimension = dimension
|
||
|
|
self.max_pixels = max_pixels
|
||
|
|
self.min_pixels = min_pixels
|
||
|
|
self.max_seq_length = max_seq_length
|
||
|
|
|
||
|
|
# Handle args
|
||
|
|
model_kwargs = model_args or {}
|
||
|
|
model_kwargs.update(kwargs)
|
||
|
|
|
||
|
|
processor_kwargs = processor_args or {}
|
||
|
|
processor_kwargs.update({
|
||
|
|
'min_pixels': min_pixels,
|
||
|
|
'max_pixels': max_pixels,
|
||
|
|
'cache_dir': cache_dir
|
||
|
|
})
|
||
|
|
|
||
|
|
# Initialize model
|
||
|
|
self.model = Qwen2VLForConditionalGeneration.from_pretrained(
|
||
|
|
model_name_or_path,
|
||
|
|
cache_dir=cache_dir,
|
||
|
|
**model_kwargs
|
||
|
|
).eval()
|
||
|
|
|
||
|
|
# Initialize processor
|
||
|
|
self.processor = AutoProcessor.from_pretrained(
|
||
|
|
processor_name_or_path or model_name_or_path,
|
||
|
|
**processor_kwargs
|
||
|
|
)
|
||
|
|
|
||
|
|
# Set padding sides
|
||
|
|
self.model.padding_side = "left"
|
||
|
|
self.processor.tokenizer.padding_side = "left"
|
||
|
|
|
||
|
|
# Store prompts
|
||
|
|
self.document_prompt = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>What is shown in this image?<|im_end|>\n<|endoftext|>"
|
||
|
|
self.query_prompt = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Query: %s<|im_end|>\n<|endoftext|>"
|
||
|
|
|
||
|
|
# Try to infer max_seq_length if not provided
|
||
|
|
if self.max_seq_length is None:
|
||
|
|
if (
|
||
|
|
hasattr(self.model, 'config')
|
||
|
|
and hasattr(self.model.config, 'max_position_embeddings')
|
||
|
|
and hasattr(self.processor.tokenizer, 'model_max_length')
|
||
|
|
):
|
||
|
|
self.max_seq_length = min(
|
||
|
|
self.model.config.max_position_embeddings,
|
||
|
|
self.processor.tokenizer.model_max_length,
|
||
|
|
)
|
||
|
|
|
||
|
|
def _smart_resize(self, height: int, width: int) -> tuple[int, int]:
|
||
|
|
h_bar = max(28, self._round_by_factor(height, 28))
|
||
|
|
w_bar = max(28, self._round_by_factor(width, 28))
|
||
|
|
if h_bar * w_bar > self.max_pixels:
|
||
|
|
beta = math.sqrt((height * width) / self.max_pixels)
|
||
|
|
h_bar = self._floor_by_factor(height / beta, 28)
|
||
|
|
w_bar = self._floor_by_factor(width / beta, 28)
|
||
|
|
elif h_bar * w_bar < self.min_pixels:
|
||
|
|
beta = math.sqrt(self.min_pixels / (height * width))
|
||
|
|
h_bar = self._ceil_by_factor(height * beta, 28)
|
||
|
|
w_bar = self._ceil_by_factor(width * beta, 28)
|
||
|
|
return w_bar, h_bar
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _round_by_factor(number: float, factor: int) -> int:
|
||
|
|
return round(number / factor) * factor
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _ceil_by_factor(number: float, factor: int) -> int:
|
||
|
|
return math.ceil(number / factor) * factor
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _floor_by_factor(number: float, factor: int) -> int:
|
||
|
|
return math.floor(number / factor) * factor
|
||
|
|
|
||
|
|
def _resize_image(self, image: Image.Image) -> Image.Image:
|
||
|
|
new_size = self._smart_resize(image.height, image.width)
|
||
|
|
return image.resize(new_size)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _decode_data_image(data_image_str: str) -> Image.Image:
|
||
|
|
header, data = data_image_str.split(',', 1)
|
||
|
|
image_data = base64.b64decode(data)
|
||
|
|
return Image.open(BytesIO(image_data))
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _is_valid_url(url: str) -> bool:
|
||
|
|
try:
|
||
|
|
result = urlparse(url)
|
||
|
|
# Check if scheme and netloc are present and scheme is http/https
|
||
|
|
return all([result.scheme in ('http', 'https'), result.netloc])
|
||
|
|
except Exception:
|
||
|
|
return False
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _is_safe_path(path: str) -> bool:
|
||
|
|
try:
|
||
|
|
# Convert to absolute path and normalize
|
||
|
|
abs_path = os.path.abspath(os.path.normpath(path))
|
||
|
|
# Check if file exists and is a regular file (not a directory or special file)
|
||
|
|
return os.path.isfile(abs_path)
|
||
|
|
except Exception:
|
||
|
|
return False
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _load_image_from_url(url: str) -> Image.Image:
|
||
|
|
try:
|
||
|
|
response = requests.get(
|
||
|
|
url,
|
||
|
|
stream=True,
|
||
|
|
timeout=10, # Add timeout
|
||
|
|
headers={'User-Agent': 'Mozilla/5.0'} # Add user agent
|
||
|
|
)
|
||
|
|
response.raise_for_status()
|
||
|
|
|
||
|
|
# Check content type
|
||
|
|
content_type = response.headers.get('content-type', '')
|
||
|
|
if not content_type.startswith('image/'):
|
||
|
|
raise ValueError(f"Invalid content type: {content_type}")
|
||
|
|
|
||
|
|
# Limit file size (e.g., 10MB)
|
||
|
|
content = BytesIO()
|
||
|
|
size = 0
|
||
|
|
max_size = 10 * 1024 * 1024 # 10MB
|
||
|
|
|
||
|
|
for chunk in response.iter_content(chunk_size=8192):
|
||
|
|
size += len(chunk)
|
||
|
|
if size > max_size:
|
||
|
|
raise ValueError("File too large")
|
||
|
|
content.write(chunk)
|
||
|
|
|
||
|
|
content.seek(0)
|
||
|
|
return Image.open(content)
|
||
|
|
except Exception as e:
|
||
|
|
raise ValueError(f"Failed to load image from URL: {str(e)}")
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _load_image_from_path(image_path: str) -> Image.Image:
|
||
|
|
try:
|
||
|
|
# Convert to absolute path and normalize
|
||
|
|
abs_path = os.path.abspath(os.path.normpath(image_path))
|
||
|
|
|
||
|
|
# Check file size before loading
|
||
|
|
file_size = os.path.getsize(abs_path)
|
||
|
|
max_size = 10 * 1024 * 1024 # 10MB
|
||
|
|
if file_size > max_size:
|
||
|
|
raise ValueError("File too large")
|
||
|
|
|
||
|
|
with Image.open(abs_path) as img:
|
||
|
|
# Make a copy to ensure file handle is closed
|
||
|
|
return img.copy()
|
||
|
|
except Exception as e:
|
||
|
|
raise ValueError(f"Failed to load image from path: {str(e)}")
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _load_image_from_bytes(image_bytes: bytes) -> Image.Image:
|
||
|
|
try:
|
||
|
|
# Check size
|
||
|
|
if len(image_bytes) > 10 * 1024 * 1024: # 10MB
|
||
|
|
raise ValueError("Image data too large")
|
||
|
|
|
||
|
|
return Image.open(BytesIO(image_bytes))
|
||
|
|
except Exception as e:
|
||
|
|
raise ValueError(f"Failed to load image from bytes: {str(e)}")
|
||
|
|
|
||
|
|
def _process_input(self, texts: List[Union[str, Image.Image, bytes]]) -> tuple[List[str], List[Image.Image]]:
|
||
|
|
processed_texts = []
|
||
|
|
processed_images = []
|
||
|
|
dummy_image = Image.new('RGB', (56, 56))
|
||
|
|
|
||
|
|
for sample in texts:
|
||
|
|
if isinstance(sample, str):
|
||
|
|
# Check if the string is a valid URL
|
||
|
|
if self._is_valid_url(sample):
|
||
|
|
try:
|
||
|
|
img = self._load_image_from_url(sample)
|
||
|
|
processed_texts.append(self.document_prompt)
|
||
|
|
processed_images.append(self._resize_image(img))
|
||
|
|
except Exception as e:
|
||
|
|
# If URL loading fails, treat as regular text
|
||
|
|
processed_texts.append(self.query_prompt % sample)
|
||
|
|
processed_images.append(dummy_image)
|
||
|
|
# Check if the string is a valid file path
|
||
|
|
elif self._is_safe_path(sample):
|
||
|
|
try:
|
||
|
|
img = self._load_image_from_path(sample)
|
||
|
|
processed_texts.append(self.document_prompt)
|
||
|
|
processed_images.append(self._resize_image(img))
|
||
|
|
except Exception as e:
|
||
|
|
# If image loading fails, treat as regular text
|
||
|
|
processed_texts.append(self.query_prompt % sample)
|
||
|
|
processed_images.append(dummy_image)
|
||
|
|
else:
|
||
|
|
# Regular text query
|
||
|
|
processed_texts.append(self.query_prompt % sample)
|
||
|
|
processed_images.append(dummy_image)
|
||
|
|
elif isinstance(sample, Image.Image):
|
||
|
|
processed_texts.append(self.document_prompt)
|
||
|
|
processed_images.append(self._resize_image(sample))
|
||
|
|
elif isinstance(sample, bytes):
|
||
|
|
try:
|
||
|
|
img = self._load_image_from_bytes(sample)
|
||
|
|
processed_texts.append(self.document_prompt)
|
||
|
|
processed_images.append(self._resize_image(img))
|
||
|
|
except Exception as e:
|
||
|
|
# If bytes can't be converted to image, use dummy
|
||
|
|
processed_texts.append(self.document_prompt)
|
||
|
|
processed_images.append(dummy_image)
|
||
|
|
|
||
|
|
return processed_texts, processed_images
|
||
|
|
|
||
|
|
def forward(self, features: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
|
||
|
|
cache_position = torch.arange(0, features['input_ids'].shape[1])
|
||
|
|
inputs = self.model.prepare_inputs_for_generation(
|
||
|
|
**features, cache_position=cache_position, use_cache=False
|
||
|
|
)
|
||
|
|
|
||
|
|
# ensure inputs are on the same device as the model
|
||
|
|
device = next(self.model.parameters()).device
|
||
|
|
inputs = {k: v.to(device) for k, v in inputs.items() if isinstance(v, torch.Tensor)}
|
||
|
|
|
||
|
|
with torch.no_grad():
|
||
|
|
output = self.model(
|
||
|
|
**inputs,
|
||
|
|
return_dict=True,
|
||
|
|
output_hidden_states=True
|
||
|
|
)
|
||
|
|
|
||
|
|
embeddings = output.hidden_states[-1][:, -1]
|
||
|
|
features['sentence_embedding'] = torch.nn.functional.normalize(
|
||
|
|
embeddings[:, :self.dimension], p=2, dim=-1
|
||
|
|
)
|
||
|
|
return features
|
||
|
|
|
||
|
|
def tokenize(self, texts: List[Union[str, Image.Image, bytes]], padding: str = 'longest') -> Dict[str, torch.Tensor]:
|
||
|
|
processed_texts, processed_images = self._process_input(texts)
|
||
|
|
|
||
|
|
return self.processor(
|
||
|
|
text=processed_texts,
|
||
|
|
images=processed_images,
|
||
|
|
videos=None,
|
||
|
|
padding=padding,
|
||
|
|
return_tensors='pt'
|
||
|
|
)
|
||
|
|
|
||
|
|
def save(self, output_path: str, safe_serialization: bool = True) -> None:
|
||
|
|
"""Save the model, tokenizer and processor to the given path."""
|
||
|
|
self.model.save_pretrained(output_path, safe_serialization=safe_serialization)
|
||
|
|
self.processor.save_pretrained(output_path)
|
||
|
|
|
||
|
|
# Save the configuration
|
||
|
|
config = {
|
||
|
|
'model_name_or_path': output_path,
|
||
|
|
'max_pixels': self.max_pixels,
|
||
|
|
'min_pixels': self.min_pixels,
|
||
|
|
'dimension': self.dimension,
|
||
|
|
'max_seq_length': self.max_seq_length,
|
||
|
|
}
|
||
|
|
|
||
|
|
config_path = os.path.join(output_path, 'sentence_bert_config.json')
|
||
|
|
with open(config_path, 'w') as f:
|
||
|
|
json.dump(config, f)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def load(input_path: str) -> 'Transformer':
|
||
|
|
"""Load a saved model from the given path."""
|
||
|
|
# Load configuration
|
||
|
|
config_path = os.path.join(input_path, 'sentence_bert_config.json')
|
||
|
|
if os.path.exists(config_path):
|
||
|
|
with open(config_path) as f:
|
||
|
|
config = json.load(f)
|
||
|
|
else:
|
||
|
|
config = {'model_name_or_path': input_path}
|
||
|
|
|
||
|
|
return Transformer(**config)
|