File size: 2,722 Bytes
3581a58 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import logging
import torch
from typing import List, Union
from tqdm import tqdm
from transformers import AutoModel, AutoTokenizer
logger = logging.getLogger(__name__)
class HFEmbedder:
def __init__(self, model_name: str = "BAAI/bge-m3", device: str = "cpu"):
self.model_name = model_name
self.device = device
logger.info(f"Loading embedding model: {model_name} on {device}")
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModel.from_pretrained(model_name).to(device)
self.model.eval()
logger.info(f"✅ Embedding model loaded successfully")
def encode(
self,
texts: Union[str, List[str]],
batch_size: int = 32,
max_length: int = 512,
normalize: bool = True,
show_progress: bool = True
) -> torch.Tensor:
"""
Encode texts into embeddings with progress bar.
Args:
texts: Single string or list of strings to encode
batch_size: Batch size for encoding
max_length: Maximum token length
normalize: Whether to L2-normalize embeddings
show_progress: Whether to show progress bar
Returns:
Tensor of shape (num_texts, embedding_dim)
"""
if isinstance(texts, str):
texts = [texts]
all_embeddings = []
# Process in batches with progress bar
iterator = range(0, len(texts), batch_size)
if show_progress:
iterator = tqdm(iterator, desc="Encoding texts", unit="batch")
for i in iterator:
batch_texts = texts[i:i + batch_size]
# Tokenize
inputs = self.tokenizer(
batch_texts,
padding=True,
truncation=True,
max_length=max_length,
return_tensors="pt"
).to(self.device)
# Get embeddings
with torch.no_grad():
outputs = self.model(**inputs)
# Use CLS token or mean pooling
embeddings = outputs.last_hidden_state[:, 0, :] # CLS token
if normalize:
embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
all_embeddings.append(embeddings.cpu())
# Concatenate all batches
final_embeddings = torch.cat(all_embeddings, dim=0)
if show_progress:
logger.info(f"✅ Encoded {len(texts)} texts -> shape {final_embeddings.shape}")
return final_embeddings |