File size: 1,584 Bytes
a7542ae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
#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) |