Update app.py

app.py

@@ -1,7 +1,7 @@
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
-import requests
+from datasets import load_dataset
 
 model_name = "Writer/palmyra-small"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -9,47 +9,18 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
 
 def get_movie_info(movie_title):
-    api_key = "20e959f0f28e6b3e3de49c50f358538a"
-    search_url = f"https://api.themoviedb.org/3/search/movie"
-    
-    # Make a search query to TMDb
-    params = {
-        "api_key": api_key,
-        "query": movie_title,
-        "language": "en-US",
-        "page": 1,
-    }
-
-    try:
-        search_response = requests.get(search_url, params=params)
-        search_data = search_response.json()
-
-        # Check if any results are found
-        if search_data.get("results"):
-            movie_id = search_data["results"][0]["id"]
-
-            # Fetch detailed information using the movie ID
-            details_url = f"https://api.themoviedb.org/3/movie/{movie_id}"
-            details_params = {
-                "api_key": api_key,
-                "language": "en-US",
-            }
+    # Load the IMDb dataset
+    imdb = load_dataset("imdb")
 
-            details_response = requests.get(details_url, params=details_params)
-            details_data = details_response.json()
+    # Search for the movie in the IMDb dataset
+    results = imdb['title'].filter(lambda x: movie_title.lower() in x.lower())
 
-            # Extract relevant information
-            title = details_data.get("title", "Unknown Title")
-            year = details_data.get("release_date", "Unknown Year")[:4]
-            genre = ", ".join(genre["name"] for genre in details_data.get("genres", []))
-
-            return f"Title: {title}, Year: {year}, Genre: {genre}"
-
-        else:
-            return "Movie not found"
-    
-    except Exception as e:
-        return f"Error: {e}"
+    # Check if any results are found
+    if len(results) > 0:
+        movie = results[0]
+        return f"Title: {movie['title']}, Year: {movie['year']}, Genre: {', '.join(movie['genre'])}"
+    else:
+        return "Movie not found"
 
 def generate_response(prompt):
     input_text_template = (
@@ -59,7 +30,7 @@ def generate_response(prompt):
         "ASSISTANT:"
     )
 
-    # Call the get_movie_info function to enrich the response
+    # Call the get_movie_info function 
     movie_info = get_movie_info(prompt)
 
     # Concatenate the movie info with the input template
@@ -83,3 +54,4 @@ def generate_response(prompt):
 iface = gr.Interface(fn=generate_response, inputs="text", outputs="text", live=True)
 iface.launch()
 
+