#ran on Kaggle
!pip install sentence-transformers
!pip install torch
import torch
from sentence_transformers import SentenceTransformer
import numpy as np
import pandas as pd
from tqdm import tqdm  # For tracking progress in batches

# Check if GPU is available
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# Load dataset
dataset = pd.read_csv('/kaggle/input/d/infamouscoder/dataset-netflix-shows/netflix_titles.csv')

# Load model to GPU if available
model = SentenceTransformer("all-MiniLM-L6-v2").to(device)

# Combine fields for embeddings
def combine_description_title_and_genre(description, listed_in, title):
    return f"{description} Genre: {listed_in} Title: {title}"

# Create combined text column
dataset['combined_text'] = dataset.apply(lambda row: combine_description_title_and_genre(row['description'], row['listed_in'], row['title']), axis=1)

# Generate embeddings in batches to save memory
batch_size = 32
embeddings = []

for i in tqdm(range(0, len(dataset), batch_size), desc="Generating Embeddings"):
    batch_texts = dataset['combined_text'][i:i+batch_size].tolist()
    batch_embeddings = model.encode(batch_texts, convert_to_tensor=True, device=device)
    embeddings.extend(batch_embeddings.cpu().numpy())  # Move to CPU to save memory

# Convert list to numpy array
embeddings = np.array(embeddings)

# Save embeddings and metadata
np.save("/kaggle/working/netflix_embeddings.npy", embeddings)
dataset[['show_id', 'title', 'description', 'listed_in']].to_csv("/kaggle/working/netflix_metadata.csv", index=False)