MohammadYaseen's picture
Update app.py
01ce151 verified
# === Install Required Libraries === #
# !pip install pandas numpy sentence-transformers faiss-cpu streamlit
# === Import Required Libraries === #
import pandas as pd
import numpy as np
from sentence_transformers import SentenceTransformer
import faiss
import streamlit as st
import os
# === Dataset Loading Function === #
def load_dataset():
"""
Provides multiple options to load the dataset: manual upload, Kaggle download, or specifying a local path.
"""
st.write("### Dataset Upload Options")
upload_option = st.radio(
"Choose how to provide the dataset:",
("Manual Upload", "Download from Kaggle", "Specify Local Path")
)
# Manual Upload
if upload_option == "Manual Upload":
st.write("#### Upload the file below:")
uploaded_file = st.file_uploader("Upload your CSV file", type="csv")
if uploaded_file is not None:
st.success("File uploaded successfully!")
return pd.read_csv(uploaded_file)
# Kaggle Download
elif upload_option == "Download from Kaggle":
st.write("#### Enter your Kaggle Dataset Path and API Key")
kaggle_dataset = st.text_input("Kaggle Dataset Path (e.g., `thedevastator/hydra-movies-dataset-directors-writers-cast-and`):")
kaggle_api_key = st.text_area("Enter your Kaggle API Key JSON content:")
if st.button("Download Dataset"):
if kaggle_dataset and kaggle_api_key:
# Set up Kaggle API
os.makedirs(os.path.expanduser("~/.kaggle"), exist_ok=True)
with open(os.path.expanduser("~/.kaggle/kaggle.json"), "w") as f:
f.write(kaggle_api_key)
os.chmod(os.path.expanduser("~/.kaggle/kaggle.json"), 0o600)
# Download dataset
os.system(f"!kaggle datasets download -d {kaggle_dataset} --unzip")
dataset_name = kaggle_dataset.split("/")[-1] + ".csv"
if os.path.exists(dataset_name):
st.success(f"Dataset {dataset_name} downloaded successfully!")
return pd.read_csv(dataset_name)
else:
st.error("Failed to download dataset. Please check your inputs.")
else:
st.warning("Please provide both the dataset path and your API key.")
# Specify Local Path
elif upload_option == "Specify Local Path":
local_path = st.text_input("Specify the full local path of your CSV file:")
if st.button("Load Dataset"):
if os.path.exists(local_path):
st.success("Dataset loaded successfully from the specified path!")
return pd.read_csv(local_path)
else:
st.error("File not found. Please check the path and try again.")
return None
# === Preprocess Data === #
def preprocess_data(df):
"""
Normalizes column names and prepares text for embeddings. Adds placeholders for missing columns if needed.
"""
# Normalize column names
df.columns = df.columns.str.strip().str.lower()
# Verify and handle missing 'genres' column
if 'genres' not in df.columns:
print("Warning: 'genres' column missing! Adding a placeholder.")
df['genres'] = "Unknown"
# Check if columns like 'title', 'summary', 'cast' exist and handle possible NaN/invalid values
df['text'] = df['title'].fillna('') + " " + df['summary'].fillna('') + " " + df['genres'] + " " + df['cast'].fillna('')
return df
# === Create Embeddings and FAISS Index === #
def create_faiss_index(df, model):
"""
Generates embeddings using a sentence-transformer model and creates a FAISS index.
"""
embeddings = model.encode(df['text'].tolist(), show_progress_bar=True)
dimension = embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(embeddings)
return index
# === Define Retrieval Function === #
def retrieve(query, model, index, df, top_k=5):
"""
Retrieves top-k results for a given query using FAISS index.
"""
query_embedding = model.encode([query])
distances, indices = index.search(query_embedding, top_k)
results = df.iloc[indices[0]].to_dict(orient="records")
return results
# === Main Function === #
if __name__ == "__main__":
# Streamlit Setup
st.title("Movie Recommendation Application with FAISS and Sentence-Transformers")
# Step 1: Load dataset
df = load_dataset()
if df is not None:
st.write("### Preview of Loaded Dataset")
st.dataframe(df.head())
# Step 2: Preprocess data
df = preprocess_data(df)
# Step 3: Create embeddings and FAISS index
st.write("### Creating Embeddings and Index...")
model = SentenceTransformer('all-MiniLM-L6-v2')
index = create_faiss_index(df, model)
# Step 4: Query Input and Result Display
query = st.text_input("Enter a movie name or keyword for recommendations:")
if query:
st.write("### Query Results")
results = retrieve(query, model, index, df)
response = ""
for i, res in enumerate(results):
response += f"**{i+1}. {res['title']} ({res['year']})**\n"
response += f"- **Genres**: {res['genres']}\n"
response += f"- **Summary**: {res['short summary']}\n"
response += f"- **Director**: {res['director']}\n"
response += f"- **Cast**: {res['cast']}\n"
response += f"- **Rating**: {res['rating']}\n\n"
st.write(response)
else:
st.write("### Please load the dataset to proceed.")