88 lines
2.9 KiB
Python
88 lines
2.9 KiB
Python
"""Geração de embeddings de texto.
|
|
|
|
Em produção usa `sentence-transformers` com um modelo multilíngue executado em
|
|
CPU. `HashEmbedder` é um embedder determinístico e sem dependências, usado nos
|
|
testes e como fallback explícito.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import logging
|
|
import math
|
|
from typing import Protocol, Sequence
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def normalize(vector: Sequence[float]) -> list[float]:
|
|
"""Normaliza um vetor para comprimento 1. Vetor nulo vira vetor nulo."""
|
|
magnitude = math.sqrt(sum(component * component for component in vector))
|
|
if magnitude == 0:
|
|
return [0.0] * len(vector)
|
|
return [component / magnitude for component in vector]
|
|
|
|
|
|
class Embedder(Protocol):
|
|
"""Interface de geração de embeddings."""
|
|
|
|
dimensions: int
|
|
|
|
def encode(self, texts: Sequence[str]) -> list[list[float]]:
|
|
"""Devolve um vetor normalizado por texto."""
|
|
...
|
|
|
|
|
|
class HashEmbedder:
|
|
"""Embedder determinístico baseado em hash.
|
|
|
|
Não captura semântica, mas serve para testes e para verificações de
|
|
integridade do pipeline (ids, contagens, filtros de metadados).
|
|
"""
|
|
|
|
def __init__(self, dimensions: int = 384):
|
|
self.dimensions = dimensions
|
|
|
|
def encode(self, texts: Sequence[str]) -> list[list[float]]:
|
|
vectors: list[list[float]] = []
|
|
for text in texts:
|
|
vector = [0.0] * self.dimensions
|
|
tokens = _tokenize(text)
|
|
for token in tokens:
|
|
digest = hashlib.sha256(token.encode("utf-8")).digest()
|
|
index = int.from_bytes(digest[:4], "big") % self.dimensions
|
|
sign = 1.0 if digest[4] % 2 == 0 else -1.0
|
|
vector[index] += sign
|
|
vectors.append(normalize(vector))
|
|
return vectors
|
|
|
|
|
|
class SentenceTransformerEmbedder:
|
|
"""Embedder de produção, baseado em sentence-transformers."""
|
|
|
|
def __init__(self, model_name: str, batch_size: int = 32, device: str = "cpu"):
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
self.model_name = model_name
|
|
self.batch_size = batch_size
|
|
self._model = SentenceTransformer(model_name, device=device)
|
|
self.dimensions = int(self._model.get_sentence_embedding_dimension())
|
|
|
|
def encode(self, texts: Sequence[str]) -> list[list[float]]:
|
|
if not texts:
|
|
return []
|
|
vectors = self._model.encode(
|
|
list(texts),
|
|
batch_size=self.batch_size,
|
|
convert_to_numpy=True,
|
|
normalize_embeddings=True,
|
|
show_progress_bar=False,
|
|
)
|
|
return [vector.tolist() for vector in vectors]
|
|
|
|
|
|
def _tokenize(text: str) -> list[str]:
|
|
"""Separa o texto em tokens simples, sem dependências externas."""
|
|
cleaned = "".join(ch.lower() if ch.isalnum() else " " for ch in text)
|
|
return [token for token in cleaned.split() if token]
|