""" Build FAISS index from movie embeddings This script should be run once to create the data files needed by the API """ import os import json import numpy as np import faiss from openai import OpenAI import requests from typing import Dict, List, Optional import time import argparse from concurrent.futures import ThreadPoolExecutor, as_completed import logging from settings import get_settings import pickle # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # Checkpoint file paths CHECKPOINT_DIR = "app/data/checkpoints" MOVIE_DATA_CHECKPOINT = f"{CHECKPOINT_DIR}/movie_data.pkl" EMBEDDINGS_CHECKPOINT = f"{CHECKPOINT_DIR}/embeddings_progress.pkl" METADATA_CHECKPOINT = f"{CHECKPOINT_DIR}/metadata_progress.pkl" def save_checkpoint(data, filepath: str): """Save checkpoint data to file""" os.makedirs(os.path.dirname(filepath), exist_ok=True) with open(filepath, 'wb') as f: pickle.dump(data, f) logger.info(f"Checkpoint saved: {filepath}") def load_checkpoint(filepath: str): """Load checkpoint data from file""" if os.path.exists(filepath): with open(filepath, 'rb') as f: data = pickle.load(f) logger.info(f"Checkpoint loaded: {filepath}") return data return None def cleanup_checkpoints(): """Remove checkpoint files after successful completion""" import shutil if os.path.exists(CHECKPOINT_DIR): shutil.rmtree(CHECKPOINT_DIR) logger.info("Checkpoint files cleaned up") class TMDBClient: """Client for TMDB API with retry and backoff""" def __init__(self, api_key: str): self.api_key = api_key self.base_url = "https://api.themoviedb.org/3" self.session = requests.Session() def _make_request(self, endpoint: str, params: dict = None, max_retries: int = 3) -> Optional[dict]: """Make API request with retry and backoff""" if params is None: params = {} params['api_key'] = self.api_key url = f"{self.base_url}{endpoint}" for attempt in range(max_retries): try: response = self.session.get(url, params=params, timeout=10) if response.status_code == 200: return response.json() elif response.status_code == 429: # Rate limit - wait and retry wait_time = 2 ** attempt logger.warning(f"Rate limited, waiting {wait_time}s before retry...") time.sleep(wait_time) continue elif response.status_code == 404: logger.warning(f"Resource not found: {url}") return None else: logger.error(f"API error {response.status_code}: {response.text}") except requests.exceptions.RequestException as e: logger.error(f"Request failed (attempt {attempt + 1}): {e}") if attempt < max_retries - 1: time.sleep(2 ** attempt) return None def get_popular_movies(self, max_pages: int = 100) -> List[int]: """Get movie IDs from popular movies pagination""" movie_ids = [] for page in range(1, max_pages + 1): logger.info(f"Fetching popular movies page {page}/{max_pages}") data = self._make_request("/movie/popular", {"page": page}) if not data: logger.error(f"Failed to fetch page {page}") break # Check if we've exceeded total pages if page > data.get('total_pages', 0): logger.info(f"Reached last page ({data.get('total_pages')})") break # Extract movie IDs for movie in data.get('results', []): movie_ids.append(movie['id']) # Rate limiting time.sleep(0.25) # 4 requests per second max logger.info(f"Collected {len(movie_ids)} movie IDs from {page} pages") return movie_ids def get_movie_details(self, movie_id: int) -> Optional[dict]: """Get detailed movie information""" return self._make_request(f"/movie/{movie_id}") def get_movie_credits(self, movie_id: int) -> Optional[dict]: """Get movie cast and crew""" return self._make_request(f"/movie/{movie_id}/credits") def fetch_movie_data(tmdb_client: TMDBClient, movie_ids: List[int], max_workers: int = 5) -> Dict[int, dict]: """Fetch detailed data for all movies with controlled parallelization""" movies_data = {} def fetch_single_movie(movie_id: int) -> tuple: """Fetch details and credits for a single movie""" try: # Get basic details details = tmdb_client.get_movie_details(movie_id) if not details: return movie_id, None # Get credits credits = tmdb_client.get_movie_credits(movie_id) if credits: details['credits'] = credits return movie_id, details except Exception as e: logger.error(f"Error fetching movie {movie_id}: {e}") return movie_id, None # Process movies in batches with controlled parallelization batch_size = 50 total_movies = len(movie_ids) for i in range(0, total_movies, batch_size): batch = movie_ids[i:i + batch_size] logger.info(f"Processing batch {i//batch_size + 1}/{(total_movies-1)//batch_size + 1} ({len(batch)} movies)") with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = {executor.submit(fetch_single_movie, movie_id): movie_id for movie_id in batch} for future in as_completed(futures): movie_id, movie_data = future.result() if movie_data: movies_data[movie_id] = movie_data # Sleep between batches to be respectful to API time.sleep(1) logger.info(f"Successfully fetched data for {len(movies_data)}/{total_movies} movies") return movies_data def create_composite_text(movie_data: Dict) -> str: """Create composite text for embedding from movie data""" parts = [] # Title if movie_data.get('title'): parts.append(f"Title: {movie_data['title']}") # Tagline if movie_data.get('tagline'): parts.append(f"Tagline: {movie_data['tagline']}") # Overview if movie_data.get('overview'): parts.append(f"Overview: {movie_data['overview']}") # Release date if movie_data.get('release_date'): parts.append(f"Release Date: {movie_data['release_date']}") # Original language if movie_data.get('original_language'): parts.append(f"Language: {movie_data['original_language']}") # Spoken languages if movie_data.get('spoken_languages'): languages = [lang.get('iso_639_1', '') for lang in movie_data['spoken_languages'] if lang.get('iso_639_1')] if languages: parts.append(f"Spoken Languages: {', '.join(languages)}") # Genres if movie_data.get('genres'): genres = [genre['name'] for genre in movie_data['genres']] parts.append(f"Genres: {', '.join(genres)}") # Production companies if movie_data.get('production_companies'): companies = [company['name'] for company in movie_data['production_companies']] if companies: parts.append(f"Production Companies: {', '.join(companies)}") # Production countries if movie_data.get('production_countries'): countries = [country['name'] for country in movie_data['production_countries']] if countries: parts.append(f"Production Countries: {', '.join(countries)}") # Budget (only if > 0) if movie_data.get('budget') and movie_data['budget'] > 0: parts.append(f"Budget: ${movie_data['budget']:,}") # Popularity if movie_data.get('popularity'): parts.append(f"Popularity: {movie_data['popularity']}") # Vote average if movie_data.get('vote_average'): parts.append(f"Vote Average: {movie_data['vote_average']}") # Vote count if movie_data.get('vote_count'): parts.append(f"Vote Count: {movie_data['vote_count']}") # Director(s) if movie_data.get('credits', {}).get('crew'): directors = [person['name'] for person in movie_data['credits']['crew'] if person['job'] == 'Director'] if directors: parts.append(f"Director: {', '.join(directors)}") # Top 5 cast if movie_data.get('credits', {}).get('cast'): top_cast = [person['name'] for person in movie_data['credits']['cast'][:5]] if top_cast: parts.append(f"Cast: {', '.join(top_cast)}") return " / ".join(parts) def get_embeddings_batch(texts: List[str], client: OpenAI, model: str = "text-embedding-3-small") -> List[List[float]]: """Get embeddings for a batch of texts with retry""" max_retries = 3 for attempt in range(max_retries): try: response = client.embeddings.create( input=texts, model=model ) return [item.embedding for item in response.data] except Exception as e: logger.error(f"Error getting embeddings (attempt {attempt + 1}): {e}") if attempt < max_retries - 1: time.sleep(2 ** attempt) else: raise def build_index(max_pages: int = 10, model: str = "text-embedding-3-small", use_faiss: bool = True): """Main function to build the FAISS index and data files""" settings = get_settings() # Initialize clients tmdb_client = TMDBClient(settings.tmdb_api_key) openai_client = OpenAI(api_key=settings.openai_api_key) # Create data directory os.makedirs("app/data", exist_ok=True) # Check for existing movie data checkpoint movies_data = load_checkpoint(MOVIE_DATA_CHECKPOINT) if movies_data is not None: logger.info(f"🔄 Resuming from checkpoint: {len(movies_data)} movies data found") else: # Step 1: Get movie IDs logger.info(f"Fetching movie IDs from TMDB (max {max_pages} pages)...") movie_ids = tmdb_client.get_popular_movies(max_pages=max_pages) if not movie_ids: logger.error("❌ No movie IDs retrieved from TMDB") return # Step 2: Fetch detailed movie data logger.info(f"Fetching detailed data for {len(movie_ids)} movies...") movies_data = fetch_movie_data(tmdb_client, movie_ids) if not movies_data: logger.error("❌ No movie data retrieved") return # Save movie data checkpoint save_checkpoint(movies_data, MOVIE_DATA_CHECKPOINT) # Step 3: Create composite texts and process embeddings in batches logger.info("Creating embeddings...") embeddings = [] id_map = {} movie_metadata = {} processed_movie_ids = set() batch_size = 20 # Process 20 texts at a time # Check for existing embedding progress embedding_checkpoint = load_checkpoint(EMBEDDINGS_CHECKPOINT) metadata_checkpoint = load_checkpoint(METADATA_CHECKPOINT) if embedding_checkpoint is not None and metadata_checkpoint is not None: embeddings = embedding_checkpoint['embeddings'] id_map = embedding_checkpoint['id_map'] processed_movie_ids = set(embedding_checkpoint['processed_movie_ids']) movie_metadata = metadata_checkpoint logger.info(f"🔄 Resuming embeddings from checkpoint: {len(embeddings)} embeddings found") else: logger.info("Starting embeddings from scratch") # Process remaining movies remaining_movies = {k: v for k, v in movies_data.items() if k not in processed_movie_ids} logger.info(f"Processing {len(remaining_movies)} remaining movies") composite_texts = [] current_movie_ids = [] for movie_id, movie_data in remaining_movies.items(): # Create composite text composite_text = create_composite_text(movie_data) composite_texts.append(composite_text) current_movie_ids.append(movie_id) # Store metadata release_year = 0 if movie_data.get("release_date"): try: release_year = int(movie_data["release_date"][:4]) except (ValueError, IndexError): release_year = 0 movie_metadata[str(movie_id)] = { "id": movie_id, "title": movie_data.get("title", ""), "year": release_year, "poster_path": movie_data.get("poster_path"), "release_date": movie_data.get("release_date"), "genres": [g["name"] for g in movie_data.get("genres", [])] } # Process batch when full if len(composite_texts) >= batch_size: logger.info(f"Processing embedding batch ({len(embeddings)} done, {len(composite_texts)} in batch)") try: batch_embeddings = get_embeddings_batch(composite_texts, openai_client, model) embeddings.extend(batch_embeddings) # Update ID mapping and processed set for i, mid in enumerate(current_movie_ids): id_map[str(mid)] = len(id_map) processed_movie_ids.add(mid) # Save progress checkpoints embedding_data = { 'embeddings': embeddings, 'id_map': id_map, 'processed_movie_ids': list(processed_movie_ids) } save_checkpoint(embedding_data, EMBEDDINGS_CHECKPOINT) save_checkpoint(movie_metadata, METADATA_CHECKPOINT) # Clear batch composite_texts = [] current_movie_ids = [] # Sleep between batches time.sleep(0.5) except Exception as e: logger.error(f"Failed to process batch: {e}") logger.info("Progress has been saved, you can restart the script to resume") return # Process remaining texts if composite_texts: logger.info(f"Processing final embedding batch ({len(composite_texts)} texts)") try: batch_embeddings = get_embeddings_batch(composite_texts, openai_client, model) embeddings.extend(batch_embeddings) for i, mid in enumerate(current_movie_ids): id_map[str(mid)] = len(id_map) processed_movie_ids.add(mid) # Save final progress embedding_data = { 'embeddings': embeddings, 'id_map': id_map, 'processed_movie_ids': list(processed_movie_ids) } save_checkpoint(embedding_data, EMBEDDINGS_CHECKPOINT) save_checkpoint(movie_metadata, METADATA_CHECKPOINT) except Exception as e: logger.error(f"Failed to process final batch: {e}") logger.info("Progress has been saved, you can restart the script to resume") return if not embeddings: logger.error("❌ No embeddings generated") return logger.info(f"Generated {len(embeddings)} embeddings") # Step 4: Save embeddings as numpy array embeddings_array = np.array(embeddings, dtype=np.float32) np.save("app/data/movies.npy", embeddings_array) logger.info(f"Saved embeddings matrix: {embeddings_array.shape}") # Step 5: Build and save FAISS index if use_faiss: logger.info("Building FAISS index...") dimension = embeddings_array.shape[1] # Choose index type based on size if len(embeddings) < 10000: # For smaller datasets, use flat index index = faiss.IndexFlatL2(dimension) else: # For larger datasets, use IVF index nlist = min(int(np.sqrt(len(embeddings))), 1000) quantizer = faiss.IndexFlatL2(dimension) index = faiss.IndexIVFFlat(quantizer, dimension, nlist) # Train the index index.train(embeddings_array) index.add(embeddings_array) faiss.write_index(index, "app/data/faiss.index") logger.info(f"FAISS index saved (type: {type(index).__name__}, dimension: {dimension})") # Step 6: Save metadata files with open("app/data/id_map.json", "w") as f: json.dump(id_map, f) with open("app/data/movie_metadata.json", "w") as f: json.dump(movie_metadata, f) logger.info("✅ Index built successfully!") logger.info(f" - {len(embeddings)} movies indexed") logger.info(f" - Embedding model: {model}") logger.info(f" - Files saved in app/data/") logger.info(f" * movies.npy: embeddings matrix") logger.info(f" * id_map.json: TMDB ID to matrix position mapping") logger.info(f" * movie_metadata.json: movie metadata") if use_faiss: logger.info(f" * faiss.index: FAISS search index") # Cleanup checkpoints cleanup_checkpoints() # Remove the old functions that are no longer needed # create_movie_embedding and load_movie_data are replaced by the new implementation if __name__ == "__main__": parser = argparse.ArgumentParser(description="Build movie embeddings index from TMDB data") parser.add_argument("--max-pages", type=int, default=10, help="Maximum pages to fetch from TMDB popular movies (default: 10)") parser.add_argument("--model", type=str, default="text-embedding-3-small", help="OpenAI embedding model to use (default: text-embedding-3-small)") parser.add_argument("--no-faiss", action="store_true", help="Skip building FAISS index") args = parser.parse_args() build_index( max_pages=args.max_pages, model=args.model, use_faiss=not args.no_faiss )