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_index = np.argmax(similarities)
    return all_quotes[most_similar_index]

def launch(input):
    out = my_pipe(input)
    return out

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

iface.launch()