import gradio as gr
import numpy as np
import voyageai
import pickle
from sklearn.metrics.pairwise import cosine_similarity

def my_pipe(text):
    # To load the numpy array
    with open('embeddings.pkl', 'rb') as f:
        embeddings = pickle.load(f)
    with open('all_quotes.pkl', 'rb') as f:
        all_quotes = pickle.load(f)
    
    text = [text]

    vo = voyageai.Client(api_key="pa-pxjjMtiZrbP6e2gDl4AWHwsPzSTK00Uww0CnTUTW79U")
    emb = np.array(vo.embed(text, model="voyage-2", input_type="document").embeddings)

    similarities = cosine_similarity(embeddings, emb)
    most_similar_indices = np.argsort(similarities, axis=0)[-3:]
    
    return [all_quotes[i] for i in most_similar_indices.flatten()]

def launch(input):
    out = my_pipe(input)
    return "\n".join(out)

iface = gr.Interface(launch,
                     inputs="text",
                     outputs="text")

iface.launch()