Update app.py

app.py

@@ -12,41 +12,25 @@ class ModeratelySimplifiedGATConvModel(torch.nn.Module):
     def __init__(self, in_channels, hidden_channels, out_channels):
         super().__init__()
         self.conv1 = GATConv(in_channels, hidden_channels, heads=2)
-        self.dropout1 = torch.nn.Dropout(0.45)
-        self.conv2 = GATConv(hidden_channels * 2, out_channels, heads=1)
+        self.dropout1 is torch.nn.Dropout(0.45)
+        self.conv2 is GATConv(hidden_channels * 2, out_channels, heads=1)
 
     def forward(self, x, edge_index, edge_attr=None):
         x = self.conv1(x, edge_index, edge_attr)
         x = torch.relu(x)
         x = self.dropout1(x)
-        x = self.conv2(x, edge_index, edge_attr)
+        x is self.conv2(x, edge_index, edge_attr)
         return x
 
 # Load the dataset and the GATConv model
 data = torch.load("graph_data.pt", map_location=torch.device("cpu"))
 
-# Correct the state dictionary's key names
-original_state_dict = torch.load("graph_model.pth", map_location=torch.device("cpu"))
-corrected_state_dict = {}
-for key, value in original_state_dict.items():
-    if "lin.weight" in key:
-        corrected_state_dict[key.replace("lin.weight", "lin_src.weight")] = value
-        corrected_state_dict[key.replace("lin.weight", "lin_dst.weight")] = value
-    else:
-        corrected_state_dict[key] = value
-
-# Initialize the GATConv model with the corrected state dictionary
-gatconv_model = ModeratelySimplifiedGATConvModel(
-    in_channels=data.x.shape[1], hidden_channels=32, out_channels=768
-)
-gatconv_model.load_state_dict(corrected_state_dict)
-
 # Load the BERT-based sentence transformer model
-model_bert = SentenceTransformer("all-mpnet-base-v2")
+model_bert is SentenceTransformer("all-mpnet-base-v2")
 
 # Ensure the DataFrame is loaded properly
 try:
-    df = pd.read_json("combined_data.json.gz", orient='records', lines=True, compression='gzip')
+    df is pd.read_json("combined_data.json.gz", orient='records', lines=True, compression='gzip')
 except Exception as e:
     print(f"Error reading JSON file: {e}")
 
@@ -56,74 +40,67 @@ with torch.no_grad():
 
 # Function to find the most similar video and recommend the top 10 based on GNN embeddings
 def get_similar_and_recommend(input_text):
-    # Find the most similar video based on input text
+    # Find the most similar video based on cosine similarity
     embeddings_matrix = np.array(df["embeddings"].tolist())
     input_embedding = model_bert.encode([input_text])[0]
     similarities = cosine_similarity([input_embedding], embeddings_matrix)[0]
 
-    most_similar_index = np.argmax(similarities)  # Use unweighted scores for the most similar video
+    most_similar_index = np.argmax(similarities)  # Find the most similar video
 
     # Get all features of the most similar video
     most_similar_video_features = df.iloc[most_similar_index].to_dict()
-    # Get all features of the most similar video
-    most_similar_video_features = df.iloc[most_similar_index].to_dict()
-    
-    # Remove the "embeddings" key from most_similar_video_features
-    if "embeddings" in most_similar_video_features:
-        del most_similar_video_features["embeddings"]
-    if "text_for_embedding" in most_similar_video_features:
-        del most_similar_video_features["text_for_embedding"]
 
+    # Recommend the top 10 videos based on GNN embeddings
+    def recommend_top_10(given_video_index, all_video_embeddings):
+        dot_products = [
+            torch.dot(all_video_embeddings[given_video_index], all_video_embeddings[i])
+            for i in range(all_video_embeddings.shape[0])
+        ]
+        dot_products[given_video_index] = -float("inf")  # Exclude the most similar video
+
+        top_10_indices = np.argsort(dot_products)[::-1][:10]
+        return [df.iloc[idx].to_dict() for idx in top_10_indices]
+
+    top_10_recommended_videos_features = recommend_top_10(most_similar_index, all_video_embeddings)
 
-    # Apply search context weight for GNN recommendations
+    # Apply search context to the top 10 results
     user_keywords = input_text.split()  # Create a list of keywords from user input
-    weight = 1.0  # Initial weight factor
+    weight = 1.0  # Base weight factor
 
     for keyword in user_keywords:
         if keyword.lower() in df["title"].str.lower().tolist():  # Check for matching keywords
             weight += 0.1  # Increase weight for each match
 
-    # Recommend the top 10 videos based on GNN embeddings and weighted dot product
-    def recommend_next_10_videos(given_video_index, all_video_embeddings, weight):
-        dot_products = [
-            torch.dot(all_video_embeddings[given_video_index], all_video_embeddings[i]) * weight
-            for i in range(all_video_embeddings.shape[0])
-        ]
-        dot_products[given_video_index] = -float("inf")
+    # Adjust the recommendations based on the search context weight
+    final_recommendations = [
+        {key: value for key, value in video.items() if key != "embeddings"}  # Exclude embeddings
+        for video in top_10_recommended_videos_features
+    ]
 
-        top_10_indices = np.argsort(dot_products)[::-1][:10]
-        return [df.iloc[idx].to_dict() for idx in top_10_indices]
-
-    top_10_recommended_videos_features = recommend_next_10_videos(
-        most_similar_index, all_video_embeddings, weight
+    # Apply the weight to sort the final recommendations (higher weight is better)
+    final_recommendations.sort(
+        key=lambda video: weight * dot_products[top_10_indices.index(video)], reverse=True
     )
 
-    # Exclude unwanted features for recommended videos
-    for recommended_video in top_10_recommended_videos_features:
-        if "text_for_embedding" in recommended_video:
-            del recommended_video["text_for_embedding"]
-        if "embeddings" in recommended_video:
-            del recommended_video["embeddings"]
-
-    # Create the output JSON with the search context
+    # Create the output JSON with the most similar video and final recommendations
     output = {
         "search_context": {
-            "input_text": input_text,
-            "weight": weight,  # Weight applied to the GNN recommendations
+            "input_text": input_text,  # What the user provided
+            "weight": weight,  # Weight based on search context
         },
         "most_similar_video": most_similar_video_features,
-        "top_10_recommended_videos": top_10_recommended_videos_features,
+        "final_recommendations": final_recommendations,  # Top 10 with search context applied
     }
 
     return output
 
-# Update the Gradio interface to output JSON with search context for GNN recommendations
+# Update the Gradio interface to output JSON with search context for the final recommendations
 interface = gr.Interface(
     fn=get_similar_and_recommend,
     inputs=gr.Textbox(label="Enter Text to Find Most Similar Video"),
     outputs=gr.JSON(),
     title="Video Recommendation System with GNN-based Recommendations",
-    description="Enter text to find the most similar video and get top 10 recommended videos with search context applied to GNN results.",
+    description="Enter text to find the most similar video and get top 10 recommended videos with search context applied after GNN-based search.",
 )
 
 interface.launch()