Spaces:
Sleeping
Sleeping
| import pandas as pd | |
| import torch | |
| from sentence_transformers import SentenceTransformer | |
| import numpy as np | |
| def create_embeddings( | |
| input_file: str, | |
| output_file: str, | |
| model_name: str = "sentence-transformers/all-MiniLM-L6-v2" | |
| ): | |
| """ | |
| Generate and save embeddings using SentenceTransformer. | |
| Args: | |
| input_file: Path to input CSV containing Title and Description columns | |
| output_file: Path to save embeddings | |
| model_name: Name of the sentence-transformer model to use | |
| """ | |
| print("Loading data...") | |
| df = pd.read_csv(input_file) | |
| print(f"Loading model: {model_name}") | |
| embedder = SentenceTransformer(model_name) | |
| # Use GPU if available | |
| if not torch.cuda.is_available(): | |
| print("Warning: No GPU found. Using CPU.") | |
| else: | |
| print("GPU found! Using CUDA.") | |
| embedder = embedder.to('cuda') | |
| # Generate combined strings for embedding | |
| print("Preparing text for embedding...") | |
| descriptions = df.apply( | |
| lambda row: f"Title: {row['Title']}; Description: {row['Description']}", | |
| axis=1 | |
| ).tolist() | |
| # Generate embeddings | |
| print("Generating embeddings...") | |
| embeddings = embedder.encode( | |
| descriptions, | |
| show_progress_bar=True, | |
| batch_size=32 if torch.cuda.is_available() else 8 | |
| ) | |
| # Save embeddings | |
| print("Saving embeddings...") | |
| embeddings_df = pd.DataFrame( | |
| embeddings, | |
| columns=[f"dim_{i}" for i in range(embeddings.shape[1])] | |
| ) | |
| embeddings_df.to_csv(output_file, index=False) | |
| print(f"Embeddings saved to {output_file}") | |
| print(f"Shape of embeddings: {embeddings.shape}") | |
| if __name__ == "__main__": | |
| create_embeddings( | |
| input_file="packages_with_titles.csv", | |
| output_file="embeddings.csv" | |
| ) |