from transformers import AutoTokenizer, AutoModel
import torch
from typing import List

from model import PersonEmbeddings

class CustomEmbeddingPipeline:
    def __init__(self, model_id="answerdotai/ModernBERT-base"):
        # Load your base tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        # Load your PersonEmbeddings
        self.model = PersonEmbeddings(model_id)
        ckpt_path = "pytorch_model.bin"
        state_dict = torch.load(ckpt_path)
        self.model.load_state_dict(state_dict)
        self.model.eval()

    def __call__(self, text: str) -> List[float]:
        # Tokenize
        inputs = self.tokenizer([text], padding=True, truncation=True, return_tensors="pt")
        with torch.no_grad():
            emb = self.model(inputs["input_ids"], inputs["attention_mask"])
        # Return the embedding of shape (1, 1536) as a Python list
        return emb[0].tolist()

def pipeline(*args, **kwargs):
    return CustomEmbeddingPipeline()