Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
-import torch
 from sentence_transformers import SentenceTransformer, util
 import pandas as pd
+import torch
 
 # -- static dataset (10 items) --
 movies = [
@@ -20,7 +20,7 @@ movies = [
 df = pd.DataFrame(movies)
 df["title_lower"] = df["title"].str.lower()
 
-# -- load sentence transformer --
+# -- load sentence transformer model --
 MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
 model = SentenceTransformer(MODEL_NAME)
 
@@ -35,18 +35,19 @@ def recommend_movie(movie_title):
     
     idx = df.index[df["title_lower"] == movie_title][0]
     query_emb = movie_embeddings[idx]
-    
+
+    # cosine similarity
     similarity = util.cos_sim(query_emb, movie_embeddings)[0]
-    top_idxs = torch.topk(similarity, k=6).indices.tolist()  # 6 to exclude itself
-    
-    recs = []
-    for i in top_idxs:
-        if i != idx:
-            recs.append(df.iloc[i]["title"])
-    
+
+    # get top 6 indices (including the movie itself)
+    top_idxs = torch.topk(similarity, k=6).indices.tolist()
+
+    # exclude the original movie
+    recs = [df.iloc[i]["title"] for i in top_idxs if i != idx]
+
     return "\n".join(recs)
 
-# -- Gradio Interface --
+# -- Gradio interface --
 interface = gr.Interface(
     fn=recommend_movie,
     inputs=gr.Textbox(label="Enter Movie title"),