Devap001's picture
Upload app.py
e6645b1
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)