first commit

requirements.txt

@@ -0,0 +1,9 @@
+streamlit
+streamlit-aggrid
+duckdb
+pandas
+requests
+json
+torch
+transformers
+textwrap

src/env_options.py

@@ -0,0 +1,57 @@
+
+import sys
+import os
+import torch
+import transformers
+
+def check_env(colab:bool=False, use_dotenv:bool=True, dotenv_path:str=None, colab_secrets:dict=None) -> tuple:
+    # Checking versions and GPU availability:
+    print(f"Python version: {sys.version}")
+    print(f"PyTorch version: {torch.__version__}")
+    print(f"Transformers version: {transformers.__version__}")
+    if torch.cuda.is_available():
+        print(f"CUDA device: {torch.cuda.get_device_name(0)}")
+        print(f"CUDA Version: {torch.version.cuda}")
+        print(f"FlashAttention available: {torch.backends.cuda.flash_sdp_enabled()}")
+    else:
+        print("No CUDA device available")
+
+    if use_dotenv:
+        print("Retrieved token(s) from .env file")
+        from dotenv import load_dotenv
+        load_dotenv(dotenv_path) # path to your dotenv file
+        hf_token = os.getenv("HF_TOKEN")
+        hf_token_write = os.getenv("HF_TOKEN_WRITE") # Only used for updating the Reddgr dataset (privileges needed)
+        openai_api_key = openai_api_key = os.getenv("OPENAI_API_KEY")
+    elif colab:
+        hf_token = colab_secrets.get('HF_TOKEN')
+        hf_token_write = colab_secrets.get('HF_TOKEN_WRITE')
+        openai_api_key = colab_secrets.get("OPENAI_API_KEY")
+    else:
+        print("Retrieved HuggingFace token(s) from environment variables")
+        hf_token = os.environ.get("HF_TOKEN")
+        hf_token_write = os.environ.get("HF_TOKEN_WRITE")
+        openai_api_key = openai_api_key = os.getenv("OPENAI_API_KEY")
+
+    def mask_token(token, unmasked_chars=4):
+        return token[:unmasked_chars] + '*' * (len(token) - unmasked_chars*2) + token[-unmasked_chars:]
+
+    if hf_token is None:
+        print("HF_TOKEN not found in the provided .env file" if use_dotenv else "HF_TOKEN not found in the environment variables")
+    if hf_token_write is None:
+        print("HF_TOKEN_WRITE not found in the provided .env file" if use_dotenv else "HF_TOKEN_WRITE not found in the environment variables")
+    if openai_api_key is None:
+        print("OPENAI_API_KEY not found in the provided .env file" if use_dotenv else "OPENAI_API_KEY not found in the environment variables")
+
+    masked_hf_token = mask_token(hf_token) if hf_token else None
+    masked_hf_token_write = mask_token(hf_token_write) if hf_token_write else None
+    masked_openai_api_key = mask_token(openai_api_key) if openai_api_key else None
+
+    if masked_hf_token:
+        print(f"Using HuggingFace token: {masked_hf_token}")
+    if masked_hf_token_write:
+        print(f"Using HuggingFace write token: {masked_hf_token_write}")
+    if masked_openai_api_key:
+        print(f"Using OpenAI token: {masked_openai_api_key}")
+
+    return hf_token, hf_token_write, openai_api_key

src/lmsys_dataset_handler.py

@@ -0,0 +1,274 @@
+import os
+import pandas as pd
+import textwrap
+import random
+from datasets import load_dataset
+from IPython.display import display
+
+class LMSYSChat1MHandler:
+    def __init__(self, hf_token, streaming=False, verbose=True):
+        self.hf_token = hf_token
+        self.streaming = streaming
+        self.lmsys_dataset = load_dataset(
+            'lmsys/lmsys-chat-1m',
+            revision="main",
+            token=self.hf_token,
+            streaming=self.streaming
+        )
+        self.verbose = verbose
+        if verbose:
+            print(self.lmsys_dataset)
+        self.df_sample = None
+        self.df_prompts = None
+        self.unwrapped_turns_df = None
+
+        if not self.streaming and verbose:
+            print('Data is cached at:\n')
+            for file_info in self.lmsys_dataset['train'].cache_files:
+                filename = file_info['filename']
+                file_size = os.path.getsize(filename)
+                i = int((len(filename) - 41) / 2)
+                print(f"Filename: {filename[:i]}*{filename[-41:]}\nSize: {file_size} bytes")
+
+    def extract_df_sample(self, n_samples=None, conversation_ids=None):
+        """
+        Extracts a sample of conversations or specific conversations based on their conversation IDs.
+
+        Parameters:
+        - n_samples (int): Number of random samples to extract. Ignored if `conversation_ids` is provided.
+        - conversation_ids (list): List of conversation IDs to extract. If provided, this takes precedence over `n_samples`.
+
+        Returns:
+        - pd.DataFrame: A DataFrame containing the extracted conversations.
+        """
+        if conversation_ids:
+            # Filter conversations based on the provided conversation IDs
+            df_sample = self.lmsys_dataset['train'].to_pandas()
+            df_sample = df_sample[df_sample['conversation_id'].isin(conversation_ids)]
+            print(f"Retrieved {len(df_sample)} conversations based on specified IDs")
+        else:
+            # Randomly sample conversations if no IDs are provided
+            if not self.streaming:    
+                df_sample = self.lmsys_dataset['train'].to_pandas().sample(n_samples)
+                print(f"Retrieved {len(df_sample)} random conversations from lmsys/lmsys-chat-1m")
+            else:
+                # Take a sample from the streamed dataset
+                streamed_samples = []
+                for i, row in enumerate(self.lmsys_dataset['train']):
+                    streamed_samples.append(row)
+                    if i + 1 == n_samples:  # Collect only the desired number of samples
+                        break
+                # Shuffle and convert the collected samples to a Pandas DataFrame
+                random.shuffle(streamed_samples)
+                df_sample = pd.DataFrame(streamed_samples)
+                
+        self.df_sample = df_sample
+        if self.verbose and len(df_sample) > 4:
+            display(df_sample.head(2))
+            print('...')
+            display(df_sample.tail(2))
+        return df_sample
+    
+    def parquet_sampling(self, n_samples):
+        base_url = "https://huggingface.co/datasets/lmsys/lmsys-chat-1m/resolve/main/data/"
+        data_files = [
+            "train-00000-of-00006-4feeb3f83346a0e9.parquet",
+            "train-00001-of-00006-4030672591c2f478.parquet",
+            "train-00002-of-00006-1779b7cec9462180.parquet",
+            "train-00003-of-00006-2fa862bfed56af1f.parquet",
+            "train-00004-of-00006-18f4bdd50c103e71.parquet",
+            "train-00005-of-00006-fe1acc5d10a9f0e2.parquet"
+        ]
+        sample_file = random.choice(data_files)
+        print(f"Sampling from {sample_file}")
+        data_files = {"train": base_url + sample_file}
+        parquet_sample = load_dataset("parquet", data_files=data_files, split="train")
+        df_sample = parquet_sample.to_pandas().sample(n_samples)
+        print(f"Retrieved {len(df_sample)} random conversations from lmsys/lmsys-chat-1m/{sample_file}")
+        self.df_sample = df_sample
+        if self.verbose and len(df_sample) > 4:
+            display(df_sample.head(2))
+            print('...')
+            display(df_sample.tail(2))
+        return df_sample
+
+    def add_turns_to_conversations(self):
+        """
+        Adds 'turn' keys to each conversation in the 'conversation' column of the dataframe.
+        """
+        self.df_sample['conversation'] = self.df_sample['conversation'].apply(
+            lambda conv: Conversation(conv).add_turns()
+        )
+        df_with_turns = self.df_sample
+        return df_with_turns
+    
+    def unwrap_turns(self):
+        """
+        Creates a dataframe where each row corresponds to a pair of user-assistant messages in a conversation and turn.
+        The 'prompt' column contains the user's message, and the 'response' column contains the assistant's message.
+        Each row includes a 'turn_id' column, which numbers the turns uniquely per conversation.
+        """
+        paired_data = []
+        for _, row in self.df_sample.iterrows():
+            conversation_id = row['conversation_id']
+            row_data = row.to_dict()
+            row_data.pop('conversation')  # Remove the 'conversation' field as it's being unwrapped
+
+            current_prompt = None
+            turn_id = None
+
+            for message in row['conversation']:
+                if message['role'] == 'user':
+                    current_prompt = message['content']
+                    turn_id = f"{conversation_id}{message['turn']:03}"  # Create turn_id
+                elif message['role'] == 'assistant' and current_prompt is not None:
+                    # Create a new row with the user-assistant pair
+                    paired_row = {
+                        **row_data,
+                        'turn_n': message['turn'],
+                        'prompt': current_prompt,
+                        'response': message['content'],
+                    }
+                    paired_data.append(paired_row)
+                    current_prompt = None  # Reset after pairing
+
+        unwrapped_turns_df = pd.DataFrame(paired_data)
+        unwrapped_turns_df.rename(columns={"turn": "conversation_turns"}, inplace=True) # The naming in the original dataset is ambiguous
+        self.unwrapped_turns_df = unwrapped_turns_df
+        return unwrapped_turns_df
+
+    def extract_prompts(self, filter_language=None, min_char_length=20, max_char_length=500, exclusions=None):
+        """
+        Extracts user prompts from the sample dataframe, optionally filtering by language and limiting the character length.
+
+        Parameters:
+        - filter_language (list of str or None): A list of specific languages to filter prompts by. If None, no language 
+          filter is applied. Examples of valid values include ['English'], ['English', 'Portuguese'], or 
+          ['Spanish', 'French', 'German'].
+        - min_char_length (int): The minimum character length for user prompts to include. Defaults to 20.
+        - max_char_length (int): The maximum character length for user prompts to include. Defaults to 500.
+        - exclusions (str or None): Path to a text file containing phrases. Prompts containing any of these phrases
+          will be excluded from the results. If None, no exclusions are applied.
+
+        Returns:
+        - pd.DataFrame: A DataFrame containing extracted prompts with columns 'prompt' and 'language'.
+        """
+        df_sample = self.df_sample
+        if filter_language:
+            extracted_data = df_sample[df_sample['language'].isin(filter_language)].apply(
+                lambda row: [
+                    {'content': entry['content'], 'language': row['language']}
+                    for entry in row['conversation']
+                    if entry['role'] == 'user' and min_char_length <= len(entry['content']) <= max_char_length
+                ], axis=1
+            ).explode().dropna()
+        else:
+            extracted_data = df_sample.apply(
+                lambda row: [
+                    {'content': entry['content'], 'language': row['language']}
+                    for entry in row['conversation']
+                    if entry['role'] == 'user' and min_char_length <= len(entry['content']) <= max_char_length
+                ], axis=1
+            ).explode().dropna()
+
+        df_prompts = pd.DataFrame(extracted_data.tolist())
+        df_prompts.rename(columns={'content': 'prompt'}, inplace=True)
+
+        orig_length = len(df_prompts)
+        if exclusions:
+            # Excluding prompts with phrases that are repeated often in this dataset
+            with open(exclusions, 'r') as f:
+                exclusions = [line.strip() for line in f.readlines()]
+            df_prompts = df_prompts[~df_prompts['prompt'].apply(lambda x: any(exclusion in x for exclusion in exclusions))]
+            print(f"Excluded {orig_length - len(df_prompts)} prompts.")
+
+        self.df_prompts = df_prompts
+        if self.verbose and len(df_sample) > 4:
+            display(df_prompts.head(2))
+            print('...')
+            display(df_prompts.tail(2))
+        return df_prompts
+    
+    def extract_prompt_sample(self):
+        prompt_sample = self.df_prompts.sample(1)['prompt'].values[0]
+        if self.verbose:
+            wrapped_message = textwrap.fill(prompt_sample, width=120)
+            print(wrapped_message)
+        return prompt_sample
+    
+    def search_conversations(self, search_term):
+        """
+        Searches the dataset for a given string and returns a DataFrame with matching records.
+
+        Parameters:
+        - search_term (str): The string to search for in the dataset.
+
+        Returns:
+        - pd.DataFrame: A DataFrame containing conversations where the search term is found.
+        """
+        if self.streaming:
+            raise ValueError("Search is not supported in streaming mode.")
+        df = self.lmsys_dataset['train'].to_pandas()
+        # Filter rows where the search term appears in the 'conversation' column
+        matching_records = df[df['conversation'].apply(
+            lambda conv: any(search_term.lower() in message['content'].lower() for message in conv)
+        )]
+        if self.verbose:
+            print(f"Found {len(matching_records)} matching conversations for search term: '{search_term}'")
+        return matching_records
+        
+    def print_language_counts(self, df):
+        language_counts = df['language'].value_counts()
+        print("Language Record Counts:")
+        print(language_counts.to_frame('Count').reset_index().rename(columns={'index': 'Language'}))
+
+
+class Conversation:
+    def __init__(self, conversation_data):
+        """
+        Initializes the Conversation object with the conversation data.
+
+        Parameters:
+        - conversation_data (list): A list of dictionaries representing a conversation.
+        """
+        self.conversation_data = conversation_data
+
+    def add_turns(self):
+        """
+        Adds a 'turn' key to each dictionary in the conversation,
+        identifying the turn (pair of user and assistant messages).
+
+        Returns:
+        - list: The updated conversation with 'turn' keys added.
+        """
+        turn_counter = 0
+        for message in self.conversation_data:
+            if message['role'] == 'user':
+                turn_counter += 1
+            message['turn'] = turn_counter
+        return self.conversation_data
+    
+    def pretty_print(self, user_prefix, assistant_prefix, width=80):
+        """
+        Prints the conversation with specified prefixes and wrapped text.
+
+        Parameters:
+        - user_prefix (str): Prefix to prepend to user messages.
+        - assistant_prefix (str): Prefix to prepend to assistant messages.
+        - width (int): Maximum characters per line for wrapping.
+        """
+        wrapper = textwrap.TextWrapper(width=width)
+        
+        for message in self.conversation_data:
+            if message['role'] == 'user':
+                prefix = user_prefix
+            elif message['role'] == 'assistant':
+                prefix = assistant_prefix
+            else:
+                continue  # Ignore roles other than 'user' and 'assistant'
+            
+            # Split on existing newlines, wrap each line, and join back with newlines
+            wrapped_content = "\n".join(
+                wrapper.fill(line) for line in message['content'].splitlines()
+            )
+            print(f"{prefix} {wrapped_content}\n")

src/lmsys_dataset_wrapper.py

@@ -0,0 +1,328 @@
+import os
+import pandas as pd
+import textwrap
+import random
+import duckdb
+import requests
+import json
+import tempfile
+from datasets import load_dataset
+from IPython.display import display
+from collections import defaultdict
+
+class DatasetWrapper:
+    def __init__(self, hf_token, dataset_name="lmsys/lmsys-chat-1m", verbose=True, 
+                 conversations_index="json/conversations_index.json", cache_size=50, request_timeout=20):
+        self.hf_token = hf_token
+        self.dataset_name = dataset_name
+        self.headers = {"Authorization": f"Bearer {self.hf_token}"}
+        self.timeout = request_timeout
+        self.cache_size = cache_size
+        self.verbose = verbose
+        parquet_list_url = f"https://datasets-server.huggingface.co/parquet?dataset={self.dataset_name}"
+        response = self._safe_get(parquet_list_url)
+        # Extract URLs from the response JSON
+        if response is not None:
+            self.parquet_urls = [file['url'] for file in response.json()['parquet_files']]
+            if self.verbose:
+                print("\nParquet URLs:")
+                for url in self.parquet_urls:
+                    print(url)
+                    head_response = self._safe_head(url)
+                    file_size = int(head_response.headers['Content-Length'])
+                    print(f"{url.split('/')[-1]}: {file_size} bytes")
+
+        # Loading the index
+        try:
+            with open(conversations_index, "r", encoding="utf-8") as f:
+                self.conversations_index = json.load(f)
+        except (FileNotFoundError, json.JSONDecodeError):
+            print(f"Conversations index file not found or invalid. Creating a new one at {conversations_index}.")
+            # Ensure directory exists
+            os.makedirs(os.path.dirname(conversations_index), exist_ok=True)
+            self.create_conversations_index(output_index_file=conversations_index)
+            with open(conversations_index, "r", encoding="utf-8") as f:
+                self.conversations_index = json.load(f)
+
+        # Initialize active conversation and DataFrame        
+        # Read from "pkl/cached_chats.pkl" if available:
+        try:
+            self.active_df = pd.read_pickle("pkl/cached_chats.pkl")
+            print(f"Loaded {len(self.active_df)} cached chats")
+            self.active_df = self.active_df.sample(self.cache_size).reset_index(drop=True)
+        except (FileNotFoundError, ValueError):
+            self.active_df = pd.DataFrame()
+            print("No cached chats found")
+        if not self.active_df.empty:
+            try:
+                self.active_conversation = Conversation(self.active_df.iloc[0])
+            except Exception as e:
+                print(f"No conversations available: {e}")
+        else:
+            self.active_conversation = None
+
+    def _safe_get(self, url):
+        if self.timeout == 0:
+            print("Timeout is set to 0. Skipping GET request.")
+            return None
+        else:
+            try:
+                response = requests.get(url, headers=self.headers, timeout=self.timeout)
+                if response.status_code != 200:
+                    raise ValueError(f"Failed to retrieve {url}. Status code: {response.status_code}")
+                return response
+            except requests.exceptions.Timeout:
+                print(f"Timeout occurred for GET {url}. Skipping.")
+                return None
+            
+    def _safe_head(self, url):
+        if self.timeout == 0:
+            print("Timeout is set to 0. Skipping HEAD request.")
+            return None
+        try:
+            response = requests.head(url, allow_redirects=True, headers=self.headers, timeout=self.timeout)
+            return response
+        except requests.exceptions.Timeout:
+            print(f"Timeout occurred for GET {url}. Skipping.")
+            return None
+
+    def extract_sample_conversations(self, n_samples):
+        url = random.choice(self.parquet_urls)
+        print(f"Sampling conversations from {url}")
+        # Download file with auth headers using requests
+        r = self._safe_get(url)
+        if r is None:
+            print(f"Timeout occurred for GET {url}. Skipping sample extraction.")
+            return self.active_df
+        # Write the downloaded content into a temporary file
+        with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as tmp:
+            tmp.write(r.content)
+            # tmp.flush()
+            tmp_path = tmp.name
+        try:
+            query_result = duckdb.query(f"SELECT * FROM read_parquet('{tmp_path}') USING SAMPLE {n_samples}").df()
+            self.active_df = query_result
+            try:
+                self.active_conversation = Conversation(query_result.iloc[0])
+            except Exception as e:
+                print(f"No conversations available: {e}")
+        finally:
+            # Clean up the temporary file
+            if os.path.exists(tmp_path):
+                os.unlink(tmp_path)
+
+        return query_result
+
+    def extract_conversations(self, conversation_ids):
+
+        # Create a lookup table for file names -> URLs
+        file_url_map = {url.split("/")[-1]: url for url in self.parquet_urls}
+
+        # Group conversation IDs by file
+        file_to_conversations = defaultdict(list)
+        for convid in conversation_ids:
+            if convid in self.conversations_index:
+                file_to_conversations[self.conversations_index[convid]].append(convid)
+
+        result_df = pd.DataFrame()
+
+        for file_name, conv_ids in file_to_conversations.items():
+            if file_name not in file_url_map:
+                print(f"File {file_name} not found in URL list, skipping.")
+                continue
+
+            file_url = file_url_map[file_name]
+            print(f"Querying file: {file_name} for {len(conv_ids)} conversations")
+
+            try:
+                r = self._safe_get(file_url)
+                if r == None:
+                    print(f"Timeout occurred for GET {file_url}. Skipping file {file_name}.")
+                    continue
+
+                with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as tmp:
+                    tmp.write(r.content)
+                    tmp_path = tmp.name
+                try:
+                    conv_id_list = "', '".join(conv_ids)
+                    query_str = f"""
+                        SELECT * FROM read_parquet('{tmp_path}') 
+                        WHERE conversation_id IN ('{conv_id_list}')
+                    """
+                    df = duckdb.query(query_str).df()
+                finally:
+                    if os.path.exists(tmp_path):
+                        os.unlink(tmp_path)
+
+                if not df.empty:
+                    print(f"Found {len(df)} conversations in {file_name}")
+                    result_df = pd.concat([result_df, df], ignore_index=True)
+
+            except Exception as e:
+                print(f"Error processing {file_name}: {e}")
+
+        self.active_df = result_df
+        try:
+            self.active_conversation = Conversation(self.active_df.iloc[0])
+        except Exception as e:
+            print(f"No conversations available: {e}")
+
+        return result_df
+    
+    def literal_text_search(self, filter_str, min_results=1):
+        # If filter_str is empty, sample random conversations
+        if filter_str == "":
+            result_df = self.extract_sample_conversations(50)
+        urls = self.parquet_urls.copy()
+        random.shuffle(urls)
+        
+        result_df = pd.DataFrame()
+
+        for url in urls:
+            print(f"Querying file: {url}")
+            r = self._safe_get(url)
+            if r == None:
+                print(f"Timeout occurred for GET {url}. Skipping file {url}.")
+                continue
+            with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as tmp:
+                tmp.write(r.content)
+                tmp_path = tmp.name
+
+            try:
+                query_str = f"""
+                    SELECT * FROM read_parquet('{tmp_path}') 
+                    WHERE contains(lower(cast(conversation as VARCHAR)), lower('{filter_str}'))
+                    """
+                df = duckdb.query(query_str).df()
+            finally:
+                if os.path.exists(tmp_path):
+                    os.unlink(tmp_path)
+
+            print(f"Found {len(df)} result(s) in {url.split('/')[-1]}")
+            
+            if len(df) > 0:
+                result_df = pd.concat([result_df, df], ignore_index=True)
+                
+            if len(result_df) >= min_results:
+                break
+        if len(result_df) == 0:
+            print("No results found. Returning empty DataFrame.")
+            placeholder_row = {'conversation_id': "No result found",
+                               'model': "-",
+                               'conversation': [
+                                {'content': '-', 'role': 'user'},
+                                {'content': '-', 'role': 'assistant'}
+                               ],
+                               'turn': "-",
+                               'language': "-",
+                               'openai_moderation': "[{'-': '-', '-': '-'}]",
+                               'redacted': "-",}
+            result_df = pd.DataFrame([placeholder_row])
+            print(result_df)
+        self.active_df = result_df
+        try:
+            self.active_conversation = Conversation(self.active_df.iloc[0])
+        except Exception as e:
+            print(f"No conversations available: {e}")
+        return result_df
+    
+    def create_conversations_index(self, output_index_file="json/conversations_index.json"):
+        """
+        Builds an index of conversation IDs from a list of Parquet file URLs.
+        Stores the index as a JSON mapping conversation IDs to their respective file names.
+        """
+        index = {}
+
+        for url in self.parquet_urls:
+            file_name = url.split('/')[-1]  # Extract file name from URL
+            print(f"Indexing file: {file_name}")
+
+            try:
+                # Download the file temporarily
+                r = requests.get(url, headers=self.headers)
+                with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as tmp:
+                    tmp.write(r.content)
+                    # tmp.flush()
+                    tmp_path = tmp.name
+                try:
+                    query = f"SELECT conversation_id FROM read_parquet('{tmp_path}')"
+                    df = duckdb.query(query).to_df()
+                finally:
+                    if os.path.exists(tmp_path):
+                        os.unlink(tmp_path)
+
+                # Map conversation IDs to file name (not the full URL)
+                for _, row in df.iterrows():
+                    index[row["conversation_id"]] = file_name
+
+            except Exception as e:
+                print(f"Error indexing {file_name}: {e}")
+
+        # Save index for fast lookup
+        with open(output_index_file, "w", encoding="utf-8") as f:
+            json.dump(index, f, indent=2)
+
+        return output_index_file
+
+
+class Conversation:
+    def __init__(self, data):
+        """
+        Initialize a conversation object either from conversation data directly or from a DataFrame row.
+        
+        Parameters:
+        - data: Can be either a list of conversation messages or a pandas Series/dict containing conversation data
+        """
+        # Handle both direct conversation data and DataFrame row
+        if isinstance(data, (pd.Series, dict)):
+            # Store all metadata separately
+            self.conversation_metadata = {}
+            for key, value in (data.items() if isinstance(data, pd.Series) else data.items()):
+                if key == 'conversation':
+                    self.conversation_data = value
+                else:
+                    self.conversation_metadata[key] = value
+        else:
+            # Direct initialization with conversation data
+            self.conversation_data = data
+            self.conversation_metadata = {}
+
+    def add_turns(self):
+        """
+        Adds a 'turn' key to each dictionary in the conversation,
+        identifying the turn (pair of user and assistant messages).
+
+        Returns:
+        - list: The updated conversation with 'turn' keys added.
+        """
+        turn_counter = 0
+        for message in self.conversation_data:
+            if message['role'] == 'user':
+                turn_counter += 1
+            message['turn'] = turn_counter
+        return self.conversation_data
+    
+    def pretty_print(self, user_prefix, assistant_prefix, width=80):
+        """
+        Prints the conversation with specified prefixes and wrapped text.
+
+        Parameters:
+        - user_prefix (str): Prefix to prepend to user messages.
+        - assistant_prefix (str): Prefix to prepend to assistant messages.
+        - width (int): Maximum characters per line for wrapping.
+        """
+        wrapper = textwrap.TextWrapper(width=width)
+        
+        for message in self.conversation_data:
+            if message['role'] == 'user':
+                prefix = user_prefix
+            elif message['role'] == 'assistant':
+                prefix = assistant_prefix
+            else:
+                continue  # Ignore roles other than 'user' and 'assistant'
+            
+            # Split on existing newlines, wrap each line, and join back with newlines
+            wrapped_content = "\n".join(
+                wrapper.fill(line) for line in message['content'].splitlines()
+            )
+            print(f"{prefix} {wrapped_content}\n")

src/text_classification_functions.py

@@ -0,0 +1,401 @@
+from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
+from datasets import Dataset
+from tqdm import tqdm
+import torch
+import numpy as np
+import os
+from langdetect import detect
+from sklearn.metrics import accuracy_score, f1_score, log_loss, confusion_matrix, ConfusionMatrixDisplay
+import matplotlib.pyplot as plt
+
+class Classifier:
+    def __init__(self, model_path, label_map, verbose = False):
+        self.model_path = model_path
+        self.classifier = pipeline("text-classification", model=model_path, tokenizer=model_path, device=0 if torch.cuda.is_available() else -1)
+        self.label_map = label_map
+        if verbose: 
+            self.print_device_information()
+    
+    def print_device_information(self):
+        # Check device information
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        device_properties = torch.cuda.get_device_properties(0) if device.type == "cuda" else "CPU Device"
+
+        print(f"Using device: {device}")
+        if device.type == "cuda":
+            print(f"Device Name: {device_properties.name}")
+            # print(f"Compute Capability: {device_properties.major}.{device_properties.minor}")
+            print(f"Total Memory: {device_properties.total_memory / 1e9:.2f} GB")
+
+    def tokenize_and_trim(self, text):
+        max_length = self.classifier.tokenizer.model_max_length
+        inputs = self.classifier.tokenizer(text, truncation=True, max_length=max_length, return_tensors="tf")
+        return self.classifier.tokenizer.decode(inputs['input_ids'][0], skip_special_tokens=True)
+
+
+    def classify_dataframe_column(self, df, target_column, feature_suffix):
+
+        tqdm.pandas()
+        df[f'trimmed_{target_column}'] = df[target_column].progress_apply(self.tokenize_and_trim)
+
+        results = []
+        for text in tqdm(df[f'trimmed_{target_column}'].tolist(), desc="Classifying"):
+            result = self.classifier(text)
+            results.append(result[0])
+
+        df[f'pred_label_{feature_suffix}'] = [self.label_map[int(result['label'].split('_')[-1])] for result in results]
+        df[f'prob_{feature_suffix}'] = [result['score'] for result in results]
+        df.drop(columns=[f'trimmed_{target_column}'], inplace=True)
+        return df
+    
+    def test_model_predictions(self, df, target_column):
+        """
+        Tests model predictions on a given dataframe column and computes evaluation metrics.
+
+        Args:
+            df (pd.DataFrame): Input dataframe containing the data.
+            target_column (str): The name of the column to classify.
+
+        Requirements:
+            - The dataframe must include a 'label' column for comparison with predictions.
+
+        Returns:
+            dict: A dictionary containing accuracy, F1 score, cross-entropy loss, 
+                and the confusion matrix.
+        """
+        # Convert pandas dataframe to Dataset
+        dataset = Dataset.from_pandas(df)
+
+        # Define a processing function for tokenization and classification
+        def process_data(batch):
+            trimmed_text = self.tokenize_and_trim(batch[target_column])
+            result = self.classifier(trimmed_text)
+            score = result[0]['score']
+            label = result[0]['label']
+            return {
+                'trimmed_text': trimmed_text,
+                'predicted_prob_0': score if label == 'LABEL_0' else 1 - score,
+                'predicted_prob_1': 1 - score if label == 'LABEL_0' else score,
+            }
+
+        # Apply processing with map
+        processed_dataset = dataset.map(process_data, batched=False)
+
+        # Convert back to pandas dataframe
+        processed_df = processed_dataset.to_pandas()
+
+        # Extract predicted probabilities and true labels
+        predicted_probs = processed_df[['predicted_prob_0', 'predicted_prob_1']].values
+        true_labels = df['label'].values
+
+        # Calculate metrics
+        accuracy = accuracy_score(true_labels, np.argmax(predicted_probs, axis=1))
+        f1 = f1_score(true_labels, np.argmax(predicted_probs, axis=1), average='weighted')
+        cross_entropy_loss = log_loss(true_labels, predicted_probs)
+
+        # Print metrics
+        print(f"Accuracy: {accuracy:.4f}")
+        print(f"F1 Score: {f1:.4f}")
+        print(f"Cross Entropy Loss: {cross_entropy_loss:.4f}")
+
+        # Confusion matrix
+        cm = confusion_matrix(true_labels, np.argmax(predicted_probs, axis=1))
+        cmap = plt.cm.Blues
+        disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=[0, 1])
+        disp.plot(cmap=cmap)
+        plt.show()
+
+        # Return metrics and probabilities for further inspection
+        return {
+            "accuracy": accuracy,
+            "f1_score": f1,
+            "cross_entropy_loss": cross_entropy_loss,
+            "confusion_matrix": cm,
+            "predicted_probs": predicted_probs  # Include reconstructed probabilities
+        }
+    
+    
+class LanguageDetector:
+    def __init__(self, dataframe):
+        """
+        Initializes the LanguageDetector with the provided DataFrame.
+        """
+        self.dataframe = dataframe
+
+    def detect_language_dataframe_column(self, target_column):
+        """
+        Detects the language of text in the specified column using langdetect and adds 
+        a 'detected_language' column to the DataFrame.
+        """
+        def detect_language(text):
+            try:
+                return detect(text)
+            except Exception:
+                return None
+
+        tqdm.pandas()
+        self.dataframe['detected_language'] = self.dataframe[target_column].progress_apply(detect_language)
+
+        return self.dataframe
+    
+
+# Classifier with Tensorflow backend
+class TensorflowClassifier(Classifier):
+    def __init__(self, model_path, label_map, verbose=False):
+        super().__init__(model_path, label_map, verbose=False)
+        self.is_tensorflow = False
+        
+        if self._is_tensorflow_model(model_path):
+            self.model = tf.keras.models.load_model(model_path)
+            self.tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")  # Adjust as per training tokenizer
+            self.is_tensorflow = True
+            if verbose:
+                print("Loaded TensorFlow model.")
+        else:
+            if verbose:
+                print("Fallback to HuggingFace pipeline.")
+
+    def _is_tensorflow_model(self, model_path):
+        return os.path.isdir(model_path) and os.path.exists(os.path.join(model_path, "saved_model.pb"))
+
+    def classify(self, text):
+        if self.is_tensorflow:
+            inputs = self.tokenizer(text, truncation=True, max_length=self.tokenizer.model_max_length, return_tensors="np")
+            logits = self.model.predict([inputs["input_ids"], inputs["attention_mask"]])
+            probabilities = tf.nn.softmax(logits).numpy()
+            label_id = np.argmax(probabilities, axis=-1).item()
+            return {
+                "label": f"LABEL_{label_id}",
+                "score": probabilities.max()
+            }
+        else:
+            return self.classifier(text)[0]
+
+    def classify_dataframe_column(self, df, target_column, feature_suffix):
+        tqdm.pandas()
+        df[f'trimmed_{target_column}'] = df[target_column].progress_apply(
+            lambda text: self.tokenizer.decode(
+                self.tokenizer(text, truncation=True, max_length=self.tokenizer.model_max_length)["input_ids"],
+                skip_special_tokens=True
+            )
+        )
+
+        if self.is_tensorflow:
+            results = [self.classify(text) for text in df[f'trimmed_{target_column}']]
+        else:
+            results = [self.classifier(text)[0] for text in df[f'trimmed_{target_column}']]
+
+        df[f'pred_label_{feature_suffix}'] = [
+            self.label_map[int(result['label'].split('_')[-1])] for result in results
+        ]
+        df[f'prob_{feature_suffix}'] = [result['score'] for result in results]
+        df.drop(columns=[f'trimmed_{target_column}'], inplace=True)
+        return df
+
+
+class ZeroShotClassifier(Classifier):
+
+    def __init__(self, model_path, tokenizer_path, candidate_labels):
+        self.model_path = model_path
+        self.candidate_labels = candidate_labels
+        self.classifier = pipeline("zero-shot-classification", model=model_path, tokenizer=tokenizer_path, clean_up_tokenization_spaces=True, device=0 if torch.cuda.is_available() else -1)
+
+    def classify_text(self, text, top_n=None, multi_label=False):
+        """
+        Classify a single text using zero-shot classification with truncated scores.
+
+        :param text: The text to classify
+        :param multi_label: Whether to allow multi-label classification
+        :return: Classification result as a dictionary with scores truncated to 3 decimals
+        """
+        classification_output = self.classifier(text, self.candidate_labels, multi_label=multi_label, clean_up_tokenization_spaces=True)
+        classification_output['scores'] = [round(score, 3) for score in classification_output['scores']]
+        if top_n is not None:
+            classification_output = {
+                'sequence': classification_output['sequence'],
+                'labels': classification_output['labels'][:top_n],
+                'scores': classification_output['scores'][:top_n]
+            }
+        return classification_output
+
+    def classify_dataframe_column(self, df, target_column, feature_suffix, multi_label=False):
+        """
+        Classify the contents of a dataframe column using zero-shot classification.
+
+        :param df: The dataframe to process
+        :param target_column: The column containing text to classify
+        :param feature_suffix: Suffix for the output columns
+        :param multi_label: Whether to allow multi-label classification
+        :return: The dataframe with classification results
+        """
+        tqdm.pandas()
+
+        # Apply the classify_text method to each row
+        results = df[target_column].progress_apply(
+            lambda text: self.classify_text(text, multi_label=multi_label)
+        )
+
+        # Extract and store results
+        df[f'top_class_{feature_suffix}'] = results.apply(lambda res: res['labels'][0])
+        df[f'top_score_{feature_suffix}'] = results.apply(lambda res: res['scores'][0])
+        df[f'full_results_{feature_suffix}'] = results.apply(lambda res: list(zip(res['labels'], res['scores'])))
+
+        return df
+    
+    def test_zs_predictions(self, df, target_column='text', true_classes_column='category', plot_conf_matrix=True):
+        """
+        Tests model predictions on a given dataset column using the zero-shot classification pipeline.
+
+        Args:
+            df (pd.DataFrame): Input dataframe containing texts for zero-shot classification.
+            target_column (str): The name of the column containing text to classify.
+            true_classes_column (str): The column containing annotated classes.
+
+        Returns:
+            dict: A dictionary containing accuracy, F1 score, and confusion matrix.
+        """
+        # Progress bar for classification
+        tqdm.pandas(desc=f"Zero-shot classification with {self.model_path}")
+        
+        # Function to classify each row
+        def classify_row(row):
+            classification_output = self.classifier(
+                row[target_column],
+                self.candidate_labels,
+                multi_label=False,
+                clean_up_tokenization_spaces=True,
+            )
+            return classification_output["labels"][0]
+
+        # Apply classification with progress bar
+        df = df.copy()
+        df.loc[:, 'predicted_class'] = df.progress_apply(classify_row, axis=1)
+        
+        # Extract true and predicted classes
+        true_classes = df[true_classes_column]
+        predicted_classes = df['predicted_class']
+
+        # Compute metrics
+        accuracy = accuracy_score(true_classes, predicted_classes)
+        f1 = f1_score(true_classes, predicted_classes, average="macro")
+        cm = confusion_matrix(true_classes, predicted_classes, labels=self.candidate_labels)
+        if plot_conf_matrix:
+            disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=self.candidate_labels)
+            fig, ax = plt.subplots(figsize=(4, 4))
+            disp.plot(cmap=plt.cm.Blues, ax=ax, colorbar=False)
+            ax.set_title(f"Zero-shot classification with {self.model_path}", fontsize=10)
+            ax.set_xlabel("Predicted label", fontsize=8) 
+            ax.set_ylabel("True label", fontsize=8)   
+
+            ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha="right", fontsize=8)
+            ax.set_yticklabels(ax.get_yticklabels(), fontsize=8)
+
+            fig.text(
+                0.5, 0.01, 
+                f"Accuracy: {accuracy:.4f} | F1 Score: {f1:.4f}",
+                ha="center",
+                fontsize=10
+            )
+            plt.tight_layout(rect=[0, 0.05, 1, 1])  # Adjust bottom margin
+            plt.show()
+
+        return {
+            "accuracy": accuracy,
+            "f1_score": f1,
+            "confusion_matrix": cm,
+            "detailed_results": df.to_dict(),  # Full dataframe with predictions
+        }
+    
+    def test_zs_predictions_with_dataset(self, df, target_column='text', true_classes_column='category', plot_conf_matrix=True):
+        dataset = Dataset.from_pandas(df)
+        def classify_text(batch):
+            classification_output = self.classifier(
+                batch[target_column],
+                self.candidate_labels,
+                multi_label=False,
+                clean_up_tokenization_spaces=True,
+            )
+            return {
+                "predicted_class": classification_output["labels"][0],
+                "predicted_scores": classification_output["scores"],
+            }
+
+        # Apply classification to the dataset
+        classified_dataset = dataset.map(classify_text, batched=False)
+        # classified_dataset = dataset.map(classify_text, batched=True, batch_size=16)
+
+        # Extract true and predicted classes
+        true_classes = classified_dataset[true_classes_column]
+        predicted_classes = classified_dataset["predicted_class"]
+
+        # Compute metrics
+        accuracy = accuracy_score(true_classes, predicted_classes)
+        f1 = f1_score(true_classes, predicted_classes, average="macro")
+
+        # Print metrics
+        print(f"Accuracy: {accuracy:.4f}")
+        print(f"F1 Score: {f1:.4f}")
+
+        # Generate confusion matrix:
+        cm = confusion_matrix(true_classes, predicted_classes, labels=self.candidate_labels)
+        if plot_conf_matrix:
+            disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=self.candidate_labels)
+            fig, ax = plt.subplots(figsize=(6, 6))  
+            disp.plot(cmap=plt.cm.Blues, ax=ax)
+            plt.xticks(rotation=45, ha="right") 
+            plt.show()
+
+        # Return metrics for further inspection
+        return {
+            "accuracy": accuracy,
+            "f1_score": f1,
+            "confusion_matrix": cm,
+            "detailed_results": classified_dataset.to_dict(), 
+        }
+    
+class MetricsComparison: 
+    def __init__(self, base_classifier, fine_tuned_classifier, base_metrics, fine_tuned_metrics):
+        self.base_classifier = base_classifier
+        self.fine_tuned_classifier = fine_tuned_classifier
+        self.base_metrics = base_metrics
+        self.fine_tuned_metrics = fine_tuned_metrics
+
+    def compare_conf_matrices(self):
+        fig, axes = plt.subplots(1, 2, figsize=(12, 6))
+        # Plot for base_classifier (left)
+        disp1 = ConfusionMatrixDisplay(confusion_matrix=self.base_metrics["confusion_matrix"], 
+                                       display_labels=self.base_classifier.candidate_labels)
+        disp1.plot(cmap=plt.cm.Blues, ax=axes[0], colorbar=False)
+        axes[0].set_title(f"Zero-shot classification with {self.base_classifier.model_path}", fontsize=10)
+        axes[0].set_xlabel("Predicted class", fontsize=8)
+        axes[0].set_ylabel("True class", fontsize=8)
+        axes[0].set_xticklabels(axes[0].get_xticklabels(), rotation=45, ha="right", fontsize=8)
+        axes[0].set_yticklabels(axes[0].get_yticklabels(), fontsize=8)
+
+        fig.text(
+            0.25, 0.01, 
+            f"Accuracy: {self.base_metrics['accuracy']:.4f} | F1 Score: {self.base_metrics['f1_score']:.4f}",
+            ha="center",
+            fontsize=10
+        )
+
+        # Plot for zs_classifier (fine-tuned) (right)
+        disp2 = ConfusionMatrixDisplay(confusion_matrix=self.fine_tuned_metrics["confusion_matrix"], 
+                                       display_labels=self.fine_tuned_classifier.candidate_labels)
+        disp2.plot(cmap=plt.cm.Blues, ax=axes[1], colorbar=False)
+        axes[1].set_title(f"ZS classification with {self.fine_tuned_classifier.model_path}", fontsize=10)
+        axes[1].set_xlabel("Predicted class", fontsize=8)
+        axes[1].set_ylabel("True class", fontsize=8)
+        axes[1].set_xticklabels(axes[1].get_xticklabels(), rotation=45, ha="right", fontsize=8)
+        axes[1].set_yticklabels(axes[1].get_yticklabels(), fontsize=8)
+
+        fig.text(
+            0.75, 0.01, 
+            f"Accuracy: {self.fine_tuned_metrics['accuracy']:.4f} | F1 Score: {self.fine_tuned_metrics['f1_score']:.4f}",
+            ha="center",
+            fontsize=10
+        )
+
+        plt.tight_layout(rect=[0, 0.05, 1, 0.95])
+        plt.show()
+