import datasets
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import gradio as gr 
from gradio.components import Label



# Load the dataset
dataset = datasets.load_dataset("SandipPalit/Movie_Dataset")
title = dataset['train']['Title']
overview = dataset['train']['Overview']

model = SentenceTransformer("sentence-transformers/all-mpnet-base-v2")

vectors = model.encode(overview) 

vector_dimension = vectors.shape[1]
index = faiss.IndexFlatL2(vector_dimension)
faiss.normalize_L2(vectors)
index.add(vectors)

def get_model_generated_vector(text): 
    search_vector = model.encode(text)
    vector = np.array([search_vector])
    faiss.normalize_L2(vector)
    return vector

def find_top_k_matched(vector):
    distances, ann = index.search(vector, k=5)
    return [title[ann[0][0]], title[ann[0][1]], title[ann[0][2]], title[ann[0][3]], title[ann[0][4]]]


def movie_recommandation(text):
    vector = get_model_generated_vector(text) 
    matches = find_top_k_matched(vector) 
#     print(matches) 
    return matches[0], matches[1], matches[2], matches[3], matches[4]

demo = gr.Interface(
    fn=movie_recommandation, 
    inputs=gr.Textbox(placeholder="Enter the Movie Name"), 
    outputs=[Label() for i in range(5)], 
    examples=[["Scarlet Macaw on Perch"], ["horror"]])

demo.launch(debug=True)