mini-russian-bert

Usage

from transformers import AutoTokenizer, AutoModel
import torch

model = AutoModel.from_pretrained("egor2014/mini-russian-bert")
tokenizer = AutoTokenizer.from_pretrained("egor2014/mini-russian-bert")

device = "cuda"

class CustomSentenceTransformer(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
        self.tokenizer = tokenizer

    def encode(self, sentences, batch_size=2, convert_to_numpy=True, **kwargs):
        all_embeddings = []
    
        for i in range(0, len(sentences), batch_size):
            batch = sentences[i:i + batch_size]  
            
            inputs = self.tokenizer(batch, return_tensors="pt", padding=True, truncation=True, max_length=512)
            inputs = {key: value.to(device) for key, value in inputs.items()}
    
            with torch.no_grad():
                outputs = self.model(**inputs, output_hidden_states=True)
            
            last_hidden_state = outputs.hidden_states[-1]
            batch_embeddings = last_hidden_state.mean(dim=1)
            all_embeddings.append(batch_embeddings.cpu())
    
        return torch.cat(all_embeddings).numpy() if convert_to_numpy else torch.cat(all_embeddings)

model.to(device)
model.eval()
custom_model = CustomSentenceTransformer(model)

custom_model.encode(["Привет, мир"])