iBrokeTheCode/Multimodal_Product_Classification

1import json2import os3 4import numpy as np5import pandas as pd6import torch7from transformers import AutoModel, AutoTokenizer8 9 10class HuggingFaceEmbeddings:11    """12    A class to handle text embedding generation using a Hugging Face pre-trained transformer model.13    This class loads the model, tokenizes the input text, generates embeddings, and provides an option14    to save the embeddings to a CSV file.15 16    Args:17        model_name (str, optional): The name of the Hugging Face pre-trained model to use for generating embeddings.18                                    Default is 'sentence-transformers/all-MiniLM-L6-v2'.19        path (str, optional): The path to the CSV file containing the text data. Default is 'data/file.csv'.20        save_path (str, optional): The directory path where the embeddings will be saved. Default is 'Models'.21        device (str, optional): The device to run the model on ('cpu' or 'cuda'). If None, it will automatically detect22                                a GPU if available; otherwise, it defaults to CPU.23 24    Attributes:25        model_name (str): The name of the Hugging Face model used for embedding generation.26        tokenizer (transformers.AutoTokenizer): The tokenizer corresponding to the chosen model.27        model (transformers.AutoModel): The pre-trained model loaded for embedding generation.28        path (str): Path to the input CSV file.29        save_path (str): Directory where the embeddings CSV will be saved.30        device (torch.device): The device on which the model and data are processed (CPU or GPU).31 32    Methods:33        get_embedding(text):34            Generates embeddings for a given text input using the pre-trained model.35 36        get_embedding_df(column, directory, file):37            Reads a CSV file, computes embeddings for a specified text column, and saves the resulting DataFrame38            with embeddings to a new CSV file in the specified directory.39 40    Example:41        embedding_instance = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2',42                                                   path='data/products.csv', save_path='output')43        text_embedding = embedding_instance.get_embedding("Sample product description.")44        embedding_instance.get_embedding_df(column='description', directory='output', file='product_embeddings.csv')45 46    Notes:47        - The Hugging Face model and tokenizer are downloaded from the Hugging Face hub.48        - The function supports large models and can run on either GPU or CPU, depending on device availability.49        - The input text will be truncated and padded to a maximum length of 512 tokens to fit into the model.50    """51 52    def __init__(53        self,54        model_name="sentence-transformers/all-MiniLM-L6-v2",55        path="data/file.csv",56        save_path=None,57        device=None,58    ):59        """60        Initializes the HuggingFaceEmbeddings class with the specified model and paths.61 62        Args:63            model_name (str, optional): The name of the Hugging Face pre-trained model. Default is 'sentence-transformers/all-MiniLM-L6-v2'.64            path (str, optional): The path to the CSV file containing text data. Default is 'data/file.csv'.65            save_path (str, optional): Directory path where the embeddings will be saved. Default is 'Models'.66            device (str, optional): Device to use for model processing. Defaults to 'cuda' if available, otherwise 'cpu'.67        """68        self.model_name = model_name69        # Load the Hugging Face tokenizer from a pre-trained model70        self.tokenizer = AutoTokenizer.from_pretrained(model_name)71 72        # Load the model from the Hugging Face model hub from the specified model name73        self.model = AutoModel.from_pretrained(model_name)74        self.path = path75        self.save_path = save_path or "Models"76 77        # Define device78        if device is None:79            # Note: If you have a mac, you may want to change 'cuda' to 'mps' to use GPU80            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")81        else:82            self.device = torch.device(device)83        print(f"Using device: {self.device}")84 85        # Move model to the specified device86        self.model.to(self.device)87        print(f"Model moved to device: {self.device}")88        print(f"Model: {model_name}")89 90    def get_embedding(self, text):91        """92        Generates embeddings for a given text using the Hugging Face model.93 94        Args:95            text (str): The input text for which embeddings will be generated.96 97        Returns:98            np.ndarray: A numpy array containing the embedding vector for the input text.99        """100        # Tokenize the input text using the Hugging Face tokenizer101        inputs = self.tokenizer(102            text, return_tensors="pt", truncation=True, padding=True, max_length=512103        )104 105        # Move the inputs to the device106        inputs = {key: value.to(self.device) for key, value in inputs.items()}107 108        with torch.no_grad():109            # Generate the embeddings using the Hugging Face model from the tokenized input110            outputs = self.model(**inputs)111 112        # Extract the embeddings from the model output, send to cpu and return the numpy array113        last_hidden_state = outputs.last_hidden_state114 115        embeddings = last_hidden_state.mean(dim=1)116        embeddings = embeddings.cpu().numpy()117 118        return embeddings[0]119 120    def get_embedding_df(self, column, directory, file):121        # Load the CSV file122        df = pd.read_csv(self.path)123        # Generate embeddings for the specified column using the `get_embedding` method124        df["embeddings"] = df[column].apply(125            lambda x: self.get_embedding(str(x)).tolist() if pd.notnull(x) else None126        )127 128        os.makedirs(directory, exist_ok=True)129 130        # Save the DataFrame with the embeddings to a new CSV file in the specified directory131        output_path = os.path.join(directory, file)132        df.to_csv(output_path, index=False)133 134        print(f"✅ Embeddings saved to {output_path}")135 136 137class GPT:138    """139    A class to interact with the OpenAI GPT API for generating text embeddings from a given dataset.140    This class provides methods to retrieve embeddings for text data and save them to a CSV file.141 142    Args:143        path (str, optional): The path to the CSV file containing the text data. Default is 'data/file.csv'.144        embedding_model (str, optional): The embedding model to use for generating text embeddings.145                                         Default is 'text-embedding-3-small'.146 147    Attributes:148        path (str): Path to the CSV file.149        embedding_model (str): The embedding model used for generating text embeddings.150 151    Methods:152        get_embedding(text):153            Generates and returns the embedding vector for the given text using the OpenAI API.154 155        get_embedding_df(column, directory, file):156            Reads a CSV file, computes the embeddings for a specified text column, and saves the embeddings157            to a new CSV file in the specified directory.158 159    Example:160        gpt_instance = GPT(path='data/products.csv', embedding_model='text-embedding-ada-002')161        text_embedding = gpt_instance.get_embedding("Sample product description.")162        gpt_instance.get_embedding_df(column='description', directory='output', file='product_embeddings.csv')163 164    Notes:165        - The OpenAI API key must be stored in a `.env` file with the variable name `OPENAI_API_KEY`.166        - The OpenAI Python package should be installed (`pip install openai`), and an active OpenAI API key is required.167    """168 169    def __init__(self, path="data/file.csv", embedding_model="text-embedding-3-small"):170        """171        Initializes the GPT class with the provided CSV file path and embedding model.172 173        Args:174            path (str, optional): The path to the CSV file containing the text data. Default is 'data/file.csv'.175            embedding_model (str, optional): The embedding model to use for generating text embeddings.176                                             Default is 'text-embedding-3-small'.177        """178        import openai179        from dotenv import find_dotenv, load_dotenv180 181        # Load the OpenAI API key from the .env file182        _ = load_dotenv(find_dotenv())  # read local .env file183        # Set the OpenAI API key184        openai.api_key = os.getenv("OPENAI_API_KEY")185 186        self.path = path187        self.embedding_model = embedding_model188 189    def get_embedding(self, text):190        """191        Generates and returns the embedding vector for the given text using the OpenAI API.192 193        Args:194            text (str): The input text to generate the embedding for.195 196        Returns:197            list: A list containing the embedding vector for the input text.198        """199        from openai import OpenAI200 201        # Instantiate the OpenAI client202        client = OpenAI()203 204        # Optional. Do text preprocessing if needed (e.g., removing newlines)205        text = text.replace("\n", " ").strip()206 207        # Call the OpenAI API to generate the embeddings and return only the embedding data208        response = client.embeddings.create(model=self.embedding_model, input=text)209 210        embeddings_np = np.array(response.data[0].embedding, dtype=np.float32)211        return embeddings_np212 213    def get_embedding_df(self, column, directory, file):214        """215        Reads a CSV file, computes the embeddings for a specified text column, and saves the results in a new CSV file.216 217        Args:218            column (str): The name of the column in the CSV file that contains the text data.219            directory (str): The directory where the output CSV file will be saved.220            file (str): The name of the output CSV file.221 222        Side Effects:223            - Saves a new CSV file containing the original data along with the computed embeddings to the specified directory.224        """225        # Load the CSV file226        df = pd.read_csv(self.path)227 228        if column not in df.columns:229            raise ValueError(f"Column '{column}' not found in CSV")230 231        # Generate embeddings in a new column 'embeddings', for the specified column using the `get_embedding` method232        df["embeddings"] = df[column].apply(233            lambda x: json.dumps(self.get_embedding(str(x)).tolist())234        )235 236        os.makedirs(directory, exist_ok=True)237 238        # Save the DataFrame with the embeddings to a new CSV file in the specified directory239        output_path = os.path.join(directory, file)240        df.to_csv(output_path, index=False)241 242        print(f"✅ Embeddings saved to {output_path}")243