import torch
from transformers import AutoModel, AutoTokenizer
from sentence_transformers import SentenceTransformer

# Load multilingual embedding model
model = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")

def encode_corpus(texts: list[str], batch_size: int = 256):
    """Encode documents into dense vector embeddings."""
    embeddings = model.encode(
        texts,
        batch_size=batch_size,
        normalize_embeddings=True,
        show_progress_bar=True,
    )
    return torch.tensor(embeddings)

# Semantic similarity search
def retrieve(query: str, corpus_emb, top_k: int = 10):
    q_emb = model.encode(query, normalize_embeddings=True)
    scores = torch.nn.functional.cosine_similarity(
        torch.tensor(q_emb).unsqueeze(0), corpus_emb
    )
    return torch.topk(scores, k=top_k)

# Fine-tune with LoRA adapters
from peft import LoraConfig, get_peft_model
config = LoraConfig(r=16, lora_alpha=32, target_modules=["q", "v"])