Similarity.py

from transformers import AutoTokenizer, AutoModel
import torch

class SimpleEmbedder:
    def __init__(self):
        print("Loading tokenizer and model...")
        self.tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
        self.model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
        print("Loaded.")

    def embed_text(self, text):
        inputs = self.tokenizer(text, return_tensors='pt')
        outputs = self.model(**inputs)
        # Mean pooling
        embeddings = outputs.last_hidden_state.mean(dim=1)
        return embeddings

embedder = SimpleEmbedder()
print(embedder.embed_text("Hello world"))