| import os
|
| import torch
|
| import json
|
| from typing import Dict, List, Union, Optional, Any
|
| from PIL import Image
|
| from transformers import AutoConfig, AutoTokenizer
|
| from custom_st import Transformer
|
|
|
| class InferenceEmbeddings:
|
| def __init__(self, model_path: str):
|
| """
|
| Initialize the embedding model for inference
|
|
|
| Args:
|
| model_path: Path to the model directory
|
| """
|
| self.model_path = model_path
|
| self.model = Transformer(
|
| model_name_or_path=model_path,
|
| model_args={"default_task": "retrieval", "trust_remote_code": True},
|
| trust_remote_code=True
|
| )
|
| self.model.eval()
|
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| self.model.to(self.device)
|
|
|
| def encode_text(self,
|
| texts: List[str],
|
| task: str = "retrieval",
|
| prompt_name: Optional[str] = None,
|
| truncate_dim: Optional[int] = None,
|
| return_multivector: bool = False,
|
| max_length: Optional[int] = None,
|
| batch_size: int = 32) -> torch.Tensor:
|
| """
|
| Encode text inputs to embeddings
|
|
|
| Args:
|
| texts: List of text inputs to encode
|
| task: Task for which to generate embeddings (retrieval, text-matching, code)
|
| prompt_name: Optional prompt type (query, passage)
|
| truncate_dim: Optional dimension to truncate embeddings to
|
| return_multivector: Whether to return multi-vector embeddings
|
| max_length: Maximum token length
|
| batch_size: Batch size for encoding
|
|
|
| Returns:
|
| Tensor of embeddings
|
| """
|
| if prompt_name:
|
|
|
| if prompt_name == "query":
|
| texts = [f"Query: {text}" for text in texts]
|
| elif prompt_name == "passage":
|
| texts = [f"Passage: {text}" for text in texts]
|
|
|
| embeddings = []
|
| for i in range(0, len(texts), batch_size):
|
| batch_texts = texts[i:i+batch_size]
|
| features = self.model.tokenize(batch_texts)
|
|
|
|
|
| for key, value in features.items():
|
| if isinstance(value, torch.Tensor):
|
| features[key] = value.to(self.device)
|
|
|
| with torch.no_grad():
|
| outputs = self.model.forward(features, task=task, truncate_dim=truncate_dim)
|
| batch_embeddings = outputs.get("sentence_embedding", None)
|
|
|
| if batch_embeddings is not None:
|
| embeddings.append(batch_embeddings.cpu())
|
|
|
| if embeddings:
|
| return torch.cat(embeddings, dim=0)
|
| else:
|
| raise RuntimeError("Failed to generate embeddings")
|
|
|
| def encode_image(self,
|
| images: List[Union[str, Image.Image]],
|
| task: str = "retrieval",
|
| truncate_dim: Optional[int] = None,
|
| return_multivector: bool = False,
|
| max_pixels: Optional[int] = None,
|
| batch_size: int = 8) -> torch.Tensor:
|
| """
|
| Encode image inputs to embeddings
|
|
|
| Args:
|
| images: List of image inputs (file paths, URLs, or PIL Images)
|
| task: Task for which to generate embeddings
|
| truncate_dim: Optional dimension to truncate embeddings to
|
| return_multivector: Whether to return multi-vector embeddings
|
| max_pixels: Maximum number of pixels for image resizing
|
| batch_size: Batch size for encoding
|
|
|
| Returns:
|
| Tensor of embeddings
|
| """
|
| embeddings = []
|
| for i in range(0, len(images), batch_size):
|
| batch_images = images[i:i+batch_size]
|
| features = self.model.tokenize(batch_images)
|
|
|
|
|
| for key, value in features.items():
|
| if isinstance(value, torch.Tensor):
|
| features[key] = value.to(self.device)
|
|
|
| with torch.no_grad():
|
| outputs = self.model.forward(features, task=task, truncate_dim=truncate_dim)
|
| batch_embeddings = outputs.get("sentence_embedding", None)
|
|
|
| if batch_embeddings is not None:
|
| embeddings.append(batch_embeddings.cpu())
|
|
|
| if embeddings:
|
| return torch.cat(embeddings, dim=0)
|
| else:
|
| raise RuntimeError("Failed to generate embeddings")
|
|
|
| def load_model(model_path: str):
|
| """
|
| Load the embedding model for inference
|
|
|
| Args:
|
| model_path: Path to the model directory
|
|
|
| Returns:
|
| Loaded model instance
|
| """
|
| return InferenceEmbeddings(model_path)
|
|
|
