The BERT Movie Search Engine is a Python-based tool that allows you to search for similar movies based on the provided query using BERT embeddings and cosine similarity. It utilizes pre-trained BERT models for computing embeddings and then compares the query with movie plot summaries to find similar movies.
Before using the BERT Movie Search Engine, make sure you have the following libraries installed:
- NumPy
- pandas
- scikit-learn (for cosine_similarity)
- transformers (for BERT model and tokenizer)
- FastAPI for web server hosting
- Uvicorn to run the API
You can install the required packages using pip
:
pip install numpy pandas scikit-learn transformers fastapi uvicorn
-
Model Setup: The movie search engine requires the
Model.py
file containing the functionsget_bert_embeddings
,tokenizer
, andmodel
for generating BERT embeddings. Make sure to import these functions into the main script for the search engine to work correctly. -
Prepare the DataFrame: Ensure that you have a CSV file containing movie data with columns "Release Year," "Title," "Genre," and "Plot." You can read the CSV file using pandas and select the relevant columns to create a DataFrame:
import pandas as pd
df = pd.read_csv("movies.csv")
df = df[['Release Year', 'Title', 'Genre', 'Plot']]
df['Index'] = df.apply(Prepare_Index, axis=1)
- Initialize the BERTSearchEngine: Create an instance of the
BERTSearchEngine
class with the prepared DataFrame:
from BERTSearchEngine import BERTSearchEngine
SearchEngineBert = BERTSearchEngine(df)
- Search for Similar Movies: Use the
Search
function to find similar movies based on a user-provided query. The function takes the search query, the initializedBERTSearchEngine
instance, and optional filters for genre and release year:
from BERTSearchEngine import Search
# Perform the search
result = Search(df, SearchEngineBert, genre="Action", year=2021)
print(result)
The search engine will return a DataFrame containing the most similar movies based on the query, sorted by similarity score in descending order.
This function takes a DataFrame (df
) and a similarity dictionary (similarity_dict
) as input. It adds a "similarity" column to the DataFrame, containing the similarity scores for each movie in the similarity_dict
. The DataFrame is then sorted based on similarity scores and returned.
This class represents the BERT-based search engine. It takes a DataFrame as input during initialization and computes BERT embeddings for the movie plot summaries. The search
method allows you to perform similarity searches based on a given query and returns a DataFrame with search results.
This function takes the DataFrame (df
), an instance of the BERTSearchEngine
(SearchEngineBert
), and optional filters for genre (genre
) and release year (year
). It prompts the user to enter a movie name as the search query, performs the similarity search, and returns a DataFrame with the most similar movies based on the query, optionally filtered by genre and release year.
from BERTSearchEngine import test
# Execute the test function
results = test()
print(results)
The test
function reads the movie data from a CSV file, initializes the BERT search engine, and performs a search for similar movies based on a user-provided query.
Please make sure to customize the file paths and data according to your specific setup.
Happy movie searching! 🎬
this documentation writed by chat GPT