from typing import Dict, List, Any
import numpy as np
from transformers import CLIPTokenizer, CLIPModel
import os


class PreTrainedPipeline():
    def __init__(self, path=""):
        # Preload all the elements you are going to need at inference.
        # For instance your model, processors, tokenizer that might be needed.
        # This function is only called once, so do all the heavy processing I/O here"""
        self.sign_ids = np.load(os.path.join(path, "sign_ids.npy"))
        self.sign_embeddings = np.load(os.path.join(path, "vanilla_large-patch14_image_embeddings.npy"))

        hf_model_path = "openai/clip-vit-large-patch14"
        self.model = CLIPModel.from_pretrained(hf_model_path)
        self.tokenizer = CLIPTokenizer.from_pretrained(hf_model_path)


    def __call__(self, inputs: str):
        """
        Args:
            inputs (:obj:`str`):
                a string to get the features from.
        Return:
            A :obj:`list` of floats: The features computed by the model.
        """
        token_inputs = self.tokenizer([inputs],  padding=True, return_tensors="pt")
        query_embed = self.model.get_text_features(**token_inputs)
        np_query_embed = query_embed.detach().cpu().numpy()[0]
        
        # Compute the cosine similarity; note the embeddings are normalized.
        cos_similarites = self.sign_embeddings @ np_query_embed
        sign_id_arg_rankings = np.argsort(cos_similarites)[::-1]
        n = 50
        top_sign_ids = self.sign_ids[sign_id_arg_rankings[:n]]
        top_sign_similarities = cos_similarites[sign_id_arg_rankings[:n]]
        return [top_sign_ids.tolist(), top_sign_similarities.tolist()]