lgcosta
8eeb6b89f4
Usa o modelo multilingual, ignora trechos curtos e o template do Llama 3.2 para a resposta sair em uma frase completa. O chat deixa a API livre durante a geração, o upload indexa o PDF e os agendamentos não são gravados duas vezes.
24 lines
816 B
Python
24 lines
816 B
Python
from sentence_transformers import SentenceTransformer
|
|
import numpy as np
|
|
|
|
_MODELS = {}
|
|
|
|
class EmbeddingService:
|
|
def __init__(self, model_name="paraphrase-multilingual-MiniLM-L12-v2"):
|
|
if model_name not in _MODELS:
|
|
_MODELS[model_name] = SentenceTransformer(model_name)
|
|
self.model = _MODELS[model_name]
|
|
self.dimension = 384
|
|
|
|
def embed(self, texts: list[str]) -> np.ndarray:
|
|
"""Generate embeddings for a list of texts."""
|
|
return self.model.encode(texts, convert_to_numpy=True)
|
|
|
|
def embed_query(self, text: str) -> np.ndarray:
|
|
"""Generate embedding for a single query."""
|
|
return self.model.encode([text], convert_to_numpy=True)[0]
|
|
|
|
def get_dimension(self) -> int:
|
|
"""Return embedding dimension."""
|
|
return self.dimension
|