can you modify this code in order to launch spacy on gpu?

please provide the whole entire corrected script.


import os
import re
import csv
import argparse
import logging
import pandas as pd
from collections import defaultdict, Counter
from typing import List, Dict, Tuple

import spacy
from gensim.models import KeyedVectors
import numpy as np
from sklearn.preprocessing import OneHotEncoder, normalize
from sklearn.cluster import AgglomerativeClustering
import hdbscan  # Ensure HDBSCAN is installed

def download_nltk_data():
    """
    Ensures that required NLTK data packages are downloaded.
    """
    import nltk
    nltk_data_packages = ['wordnet', 'averaged_perceptron_tagger', 'omw-1.4']
    for package in nltk_data_packages:
        try:
            nltk.data.find(f'corpora/{package}')
            logging.info(f"NLTK package '{package}' already exists.")
        except LookupError:
            nltk.download(package)
            logging.info(f"Downloaded NLTK package '{package}'.")

def parse_arguments():
    """
    Parses command-line arguments.
    """
    parser = argparse.ArgumentParser(description='Process and cluster keywords from a CSV file.')
    parser.add_argument('--input_file', type=str, required=True, help='Path to the input CSV file.')
    parser.add_argument('--save_lemmatized', action='store_true', help='Save lemmatized keywords and counts.')
    parser.add_argument('--frequency_threshold', type=int, default=10,
                        help='Frequency threshold for keywords. Default is 10.')
    parser.add_argument('--category_limit', type=int, default=100,
                        help='Limit of top N keywords per category. Default is 100.')
    parser.add_argument('--remove_duplicates', action='store_true',
                        help='Remove duplicates in clustered keywords per row.')
    return parser.parse_args()

def main():
    # Parse command-line arguments
    args = parse_arguments()
    input_file = args.input_file
    save_lemmatized = args.save_lemmatized
    frequency_threshold = args.frequency_threshold
    category_limit = args.category_limit
    remove_duplicates = args.remove_duplicates

    # Check if input file exists
    if not os.path.exists(input_file):
        raise FileNotFoundError(f"Input file '{input_file}' does not exist.")

    # Set up logging
    input_dir = os.path.dirname(os.path.abspath(input_file))
    input_filename = os.path.basename(input_file)
    input_name_no_ext = os.path.splitext(input_filename)[0]
    processed_dir = os.path.join(input_dir, 'processed')
    os.makedirs(processed_dir, exist_ok=True)
    log_file = os.path.join(processed_dir, f'log_{input_name_no_ext}.log')
    logging.basicConfig(filename=log_file, level=logging.INFO,
                        format='%(asctime)s %(levelname)s:%(message)s')
    # Also log to stdout
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_formatter = logging.Formatter('%(asctime)s %(levelname)s:%(message)s')
    console_handler.setFormatter(console_formatter)
    logging.getLogger().addHandler(console_handler)

    logging.info(f"Processing input file: {input_file}")

    # Download NLTK data if necessary
    download_nltk_data()

    # Load spaCy model
    try:
        nlp = spacy.load('en_core_web_trf')  # Ensure 'en_core_web_trf' is installed
        logging.info("Loaded spaCy model 'en_core_web_trf' successfully.")
    except OSError:
        logging.error("spaCy model 'en_core_web_trf' not found. Please install it using:\n"
                      "python -m spacy download en_core_web_trf")
        raise

    # Labels (customizable)
    labels = ['Type', 'Color', 'Wear Location', 'Material', 'Style']  # Modify as needed

    # Define output file paths
    output_file = os.path.join(processed_dir, f'clustered_{input_filename}')
    one_hot_encoded_file = os.path.join(processed_dir, f'onehot_{input_name_no_ext}.tsv')
    lemmatized_counts_output_file = os.path.join(processed_dir, f'lemmatized_keyword_counts.csv')

    # ------------------------------#
    # Step 1: Cleaning             #
    # ------------------------------#
    logging.info("Starting CSV cleaning...")
    df_clean = clean_csv(input_file, labels)
    logging.info("CSV cleaning completed.")

    # ------------------------------#
    # Step 2: Lemmatization        #
    # ------------------------------#
    logging.info("Starting lemmatization...")
    if save_lemmatized:
        per_row_df, keyword_counts = lemmatize_keywords(df_clean, lemmatized_counts_output_file, nlp)
    else:
        per_row_df, keyword_counts = lemmatize_keywords(df_clean, None, nlp)
    # Check if per_row_df is None
    if per_row_df is None:
        logging.error("Error in lemmatization.")
        return

    # ------------------------------#
    # Step 3: Clustering           #
    # ------------------------------#
    logging.info("Starting clustering of keywords...")
    clustered_keyword_mapping = perform_clustering(
        keyword_counts, nlp,
        frequency_threshold=frequency_threshold
    )
    logging.info("Keyword clustering completed.")

    # ------------------------------#
    # Step 4: Clustered Keywords   #
    # ------------------------------#
    per_row_df['Clustered Keywords'] = per_row_df['Lemmatized Keywords'].apply(
        lambda kws: ', '.join([clustered_keyword_mapping.get(kw, kw) for kw in kws.split(', ')])
    )

    # ------------------------------#
    # Step 5: Categorization       #
    # ------------------------------#
    logging.info("Starting categorization of clustered keywords...")
    category_keyword_counts = categorize_and_count_clustered_keywords(
        per_row_df, labels=labels, category_limit=category_limit
    )
    logging.info("Categorization completed.")

    # ------------------------------#
    # Step 6: Final Processing     #
    # ------------------------------#
    logging.info("Generating final clustered keywords and categories...")
    per_row_df = generate_clustered_keywords_and_categories(
        per_row_df, clustered_keyword_mapping, category_keyword_counts, labels=labels,
        remove_duplicates=remove_duplicates
    )
    logging.info("Final processing completed.")

    # Optionally, write category counts to CSV files
    if save_lemmatized:
        for category, counts in category_keyword_counts.items():
            category_counts_df = pd.DataFrame({'Keyword': list(counts.keys()), 'Count': list(counts.values())})
            category_counts_df = category_counts_df.sort_values(by='Count', ascending=False)
            category_counts_output_file = os.path.join(processed_dir, f'{category}_keywords_counts.csv')
            category_counts_df.to_csv(category_counts_output_file, index=False)
            logging.info(f"{category} keyword counts saved to '{category_counts_output_file}'.")

    # ------------------------------#
    # Step 7: Save Results         #
    # ------------------------------#
    per_row_df.to_csv(output_file, index=False)
    logging.info(f"Process completed. Output saved to '{output_file}'.")

    # ------------------------------#
    # Step 8: One-Hot Encoding     #
    # ------------------------------#
    logging.info("Starting one-hot encoding of categorical features...")
    encode_categorical_features(
        per_row_df, one_hot_encoded_file,
        id_column='Image Name', remove_extension='.jpg',
        labels=labels,
        sep='\t', header=False, index=False
    )
    logging.info(f"One-hot encoded features saved to '{one_hot_encoded_file}'.")


def clean_csv(input_file: str, labels: List[str]) -> pd.DataFrame:
    """
    Cleans the input CSV file and returns a DataFrame.

    Parameters:
        input_file (str): Path to the input CSV file.
        labels (List[str]): List of labels to process.

    Returns:
        pd.DataFrame: Cleaned DataFrame with 'Image Name' and 'Keywords' columns.
    """
    logging.info("Starting to clean the CSV file...")
    total_row_count = 0
    empty_row_count = 0
    na_count = 0
    no_comma_count = 0
    no_comma_rows = []
    cleaned_data = []
    # Read the entire content of the file
    try:
        with open(input_file, 'r', newline='', encoding='utf-8') as infile:
            content = infile.read()
    except Exception as e:
        logging.error(f"Error reading input file: {e}")
        raise

    # Remove newline characters inside quoted strings
    content = re.sub(r'("\[.*?")\n', r'\1 ', content)
    # Use StringIO to read the cleaned content as a CSV
    import io
    temp_file = io.StringIO(content)
    reader = csv.reader(temp_file)

    # Skip the first line (header)
    try:
        header = next(reader)
    except StopIteration:
        logging.error("Input file is empty.")
        raise ValueError("Input file is empty.")

    # Process each row
    for row in reader:
        total_row_count += 1
        if not row:
            continue
        jpg_file = row[0]
        if len(row) < 2:
            logging.warning(f"Row {total_row_count} doesn't have enough columns: {row}")
            keywords = []
        else:
            keywords = extract_keywords(row[1], labels)
        # Count occurrences of "N/A" and empty spaces
        na_count += keywords.count("N/A")
        na_count += sum(1 for k in keywords if k.strip() == "")
        # Check if the row is empty or contains only "N/A" or empty spaces
        if not keywords or all(k.strip().upper() == 'N/A' or k.strip() == '' for k in keywords):
            empty_row_count += 1
            continue  # Skip this row
        # Concatenate the keywords into a single string
        keywords_str = ', '.join(keywords)
        # Replace multiple spaces with a comma and a single space
        keywords_str = re.sub(r'\s{2,}', ', ', keywords_str)
        # Check if the row has no commas
        if ',' not in keywords_str:
            no_comma_count += 1
            no_comma_rows.append(f'{jpg_file},{keywords_str}')
        # Add to cleaned data
        cleaned_data.append({'Image Name': jpg_file, 'Keywords': keywords_str})
    # Now, create a DataFrame from the cleaned_data
    df = pd.DataFrame(cleaned_data)
    # Clean the 'Keywords' column further
    df = clean_keywords_in_df(df, labels)
    # Log statistics
    logging.info(f"Total rows processed: {total_row_count}")
    logging.info(f"Number of empty or N/A rows skipped: {empty_row_count}")
    logging.info(f"Occurrences of N/A (including empty spaces): {na_count}")
    logging.info(f"Number of rows without commas: {no_comma_count}")
    if no_comma_rows:
        logging.info("Rows without commas:")
        for row in no_comma_rows[:10]:  # Limiting to first 10 for brevity
            logging.info(row)
        if len(no_comma_rows) > 10:
            logging.info(f"...and {len(no_comma_rows) - 10} more.")
    logging.info("CSV cleaning completed successfully.")
    return df

def extract_keywords(line: str, labels: List[str]) -> List[str]:
    """
    Extracts keywords from a line.

    Parameters:
        line (str): The line from which to extract keywords.
        labels (List[str]): List of labels to be removed.

    Returns:
        List[str]: List of extracted keywords.
    """
    # First, try to extract keywords between curly braces
    keywords = re.findall(r'\{([^}]+)\}', line)
    if keywords:
        # Further split if multiple keywords are within braces
        return [kw.strip() for group in keywords for kw in group.split(',')]
    # Remove newline characters from the keywords section
    line = line.replace('\n', ' ')
    # Remove the labels using word boundaries and considering possible formats
    for label in labels:
        # Patterns like "Label:", "Label -", "[Label]:", etc.
        pattern = rf'\b{re.escape(label)}\b[: -]?'
        line = re.sub(pattern, '', line, flags=re.IGNORECASE)
    # Remove any '- ' (dash and space) using exact string replacement
    line = line.replace('- ', '')
    # Remove single quotes around keywords using regex
    line = re.sub(r"'([^']+)'", r'\1', line)
    # Remove any remaining brackets
    line = line.replace('[', '').replace(']', '').replace('{', '').replace('}', '')
    # Split keywords by commas and strip whitespace
    keywords = [keyword.strip() for keyword in line.split(',') if keyword.strip()]
    return keywords

def clean_keywords_in_df(df: pd.DataFrame, labels: List[str]) -> pd.DataFrame:
    """
    Cleans the 'Keywords' column in the DataFrame.

    Parameters:
        df (pd.DataFrame): DataFrame containing the keywords to be cleaned.
        labels (List[str]): List of labels to be removed.

    Returns:
        pd.DataFrame: DataFrame with cleaned 'Keywords' column.
    """
    # Construct a regex pattern from the labels with word boundaries
    labels_pattern = r'\b(' + '|'.join(map(re.escape, labels)) + r')\b[: -]*'
    # Replace labels with an empty string using regex
    df['Keywords'] = df['Keywords'].str.replace(labels_pattern, '', regex=True)
    # Remove any '- ' (dash and space)
    df['Keywords'] = df['Keywords'].str.replace('- ', '', regex=False)
    # Remove single quotes around keywords
    df['Keywords'] = df['Keywords'].str.replace(r"'([^']+)'", r'\1', regex=True)
    # Replace multiple spaces with a comma and a single space
    df['Keywords'] = df['Keywords'].str.replace(r'\s{2,}', ', ', regex=True)
    # Remove leading/trailing commas and spaces
    df['Keywords'] = df['Keywords'].str.strip(' ,')
    return df

def lemmatize_keywords(df: pd.DataFrame, lemmatized_counts_output_file: str, nlp) -> Tuple[pd.DataFrame, Counter]:
    """
    Lemmatizes the keywords in the DataFrame and returns a DataFrame with lemmatized keywords.
    Optionally saves the lemmatized keyword counts to a CSV file.

    Parameters:
        df (pd.DataFrame): DataFrame containing 'Image Name' and 'Keywords' columns.
        lemmatized_counts_output_file (str): Path to save the lemmatized keyword counts CSV.
        nlp: spaCy NLP model.

    Returns:
        Tuple[pd.DataFrame, Counter]: DataFrame with lemmatized keywords and keyword counts.
    """
    logging.info("Starting lemmatization of keywords with spaCy...")
    per_row_data = []
    keyword_counts = Counter()
    
    # Define a custom mapping for irregular lemmatization cases
    irregular_lemmas = {
        'feet': 'foot',
        'geese': 'goose',
        'mice': 'mouse',
        'men': 'man',
        'women': 'woman',
        'children': 'child',
        'sporty': 'sport',
        'dirty': 'dirt',
        'lives': 'life',
        'ears': 'ear',
        'legs': 'leg',
        'active wear': 'active wear',  # Assuming 'active wear' is treated as a compound noun
        # Add more irregular mappings as needed
    }

    for index, row in df.iterrows():
        image_name = row.get('Image Name', '')
        keywords_str = row.get('Keywords', '')
        if pd.isna(keywords_str):
            keywords_str = ''
        # Split keywords by comma and strip whitespace
        extracted_keywords = [kw.strip() for kw in keywords_str.split(',')]
        lemmatized_keywords = []
        for kw in extracted_keywords:
            kw = kw.strip()
            if not kw:
                continue
            doc = nlp(kw)
            lemmas = []
            for token in doc:
                lemmas.append(token.lemma_)
            lemma = ' '.join(lemmas).lower()
            # Handle irregular lemmas
            lemma = irregular_lemmas.get(lemma, lemma)
            lemmatized_keywords.append(lemma)
        # Update global counts
        keyword_counts.update(lemmatized_keywords)
        # Store per-row data
        per_row_data.append({
            'Image Name': image_name,
            'Keywords': ', '.join(extracted_keywords),
            'Lemmatized Keywords': ', '.join(lemmatized_keywords)
        })
    
    # Create a DataFrame from per_row_data
    per_row_df = pd.DataFrame(per_row_data)
    
    if lemmatized_counts_output_file:
        # Write the lemmatized keyword counts to CSV
        lemmatized_counts = [{'Lemmatized Keyword': kw, 'Count': count}
                             for kw, count in keyword_counts.items()]
        lemmatized_counts_df = pd.DataFrame(lemmatized_counts)
        # Sort by Count in descending order
        lemmatized_counts_df = lemmatized_counts_df.sort_values(by='Count', ascending=False)
        # Write to CSV
        try:
            lemmatized_counts_df.to_csv(lemmatized_counts_output_file, index=False)
            logging.info(f"Lemmatized keyword counts saved to '{lemmatized_counts_output_file}'.")
        except Exception as e:
            logging.error(f"Error saving lemmatized keyword counts: {e}")
            raise
    logging.info("Lemmatization completed.")
    return per_row_df, keyword_counts

def perform_clustering(keyword_counts: Counter, nlp, frequency_threshold: int = 10) -> Dict[str, str]:
    """
    Clusters keywords using HDBSCAN and returns a mapping from each keyword to its cluster representative.

    Parameters:
        keyword_counts (Counter): Counts of lemmatized keywords.
        nlp: spaCy NLP model.
        frequency_threshold (int): Frequency threshold to consider for clustering.

    Returns:
        Dict[str, str]: A mapping from keywords to their cluster representatives.
    """
    logging.info("Starting clustering using HDBSCAN...")
    # Separate frequent and infrequent keywords
    frequent_keywords = [kw for kw, count in keyword_counts.items() if count > frequency_threshold]
    infrequent_keywords = [kw for kw, count in keyword_counts.items() if count <= frequency_threshold]
    logging.info(f"Number of frequent keywords (> {frequency_threshold}): {len(frequent_keywords)}")
    logging.info(f"Number of infrequent keywords (<= {frequency_threshold}): {len(infrequent_keywords)}")
    
    # Keep frequent keywords as-is
    clustered = {kw: kw for kw in frequent_keywords}
    
    # Cluster infrequent keywords using HDBSCAN
    if infrequent_keywords:
        logging.info("Preparing vectors for infrequent keywords...")
        vectors = []
        valid_keywords = []
        for kw in infrequent_keywords:
            doc = nlp(kw)
            if doc.vector_norm:
                vectors.append(doc.vector)
                valid_keywords.append(kw)
            else:
                logging.warning(f"No vector found for keyword: '{kw}'. Skipping.")
        if not vectors:
            logging.warning("No valid vectors found for infrequent keywords. Skipping clustering.")
            return clustered
        vectors = np.array(vectors)
        # Normalize vectors to unit length
        vectors = normalize(vectors)
        logging.info("Clustering infrequent keywords with HDBSCAN...")
        # Initialize HDBSCAN
        clusterer = hdbscan.HDBSCAN(min_cluster_size=2, metric='euclidean', cluster_selection_method='eom')
        cluster_labels = clusterer.fit_predict(vectors)
        num_clusters = len(set(cluster_labels)) - (1 if -1 in cluster_labels else 0)
        num_noise = list(cluster_labels).count(-1)
        logging.info(f"HDBSCAN found {num_clusters} clusters and {num_noise} noise points.")
        
        # Organize keywords by clusters
        clusters = defaultdict(list)
        for label, kw in zip(cluster_labels, valid_keywords):
            clusters[label].append(kw)
        
        # Select the most frequent keyword in each cluster as the representative
        for label, kws in clusters.items():
            if label == -1:
                # Noise points remain as-is
                for kw in kws:
                    clustered[kw] = kw
            else:
                # Representative is the most frequent keyword
                most_frequent = max(kws, key=lambda w: keyword_counts[w])
                for kw in kws:
                    clustered[kw] = most_frequent
        logging.info("Clustering completed.")
    else:
        logging.info("No infrequent keywords to cluster.")
    
    return clustered

def categorize_and_count_clustered_keywords(per_row_df: pd.DataFrame, labels: List[str], category_limit: int = 100) -> Dict[str, Counter]:
    """
    Categorizes clustered keywords and counts them per category.

    Parameters:
        per_row_df (pd.DataFrame): DataFrame containing 'Clustered Keywords'.
        labels (List[str]): List of category labels.
        category_limit (int): Limit of top N keywords per category.

    Returns:
        Dict[str, Counter]: Dictionary of keyword counts per category.
    """
    logging.info("Categorizing and counting clustered keywords...")
    # Initialize counters for each category
    category_keyword_counts = {label: Counter() for label in labels}
    # Build initial mapping of positions to labels
    for idx, row in per_row_df.iterrows():
        clustered_keywords = row['Clustered Keywords'].split(', ')
        for kw in clustered_keywords:
            # For each keyword, we will not assign based on position but collect counts for potential categories
            # This step is adjusted in 'generate_clustered_keywords_and_categories'
            pass  # We will build the counts in the next function
    # Since we're changing the assignment logic, we return empty counters here
    logging.info("Categorization and counting completed.")
    return category_keyword_counts

def generate_clustered_keywords_and_categories(per_row_df: pd.DataFrame, clustered_keyword_mapping: Dict[str, str],
                                               category_keyword_counts: Dict[str, Counter], labels: List[str],
                                               remove_duplicates: bool = False) -> pd.DataFrame:
    """
    Generates the final DataFrame with clustered keywords and category assignments.

    Parameters:
        per_row_df (pd.DataFrame): DataFrame containing 'Lemmatized Keywords'.
        clustered_keyword_mapping (Dict[str, str]): Mapping of keywords to their cluster representatives.
        category_keyword_counts (Dict[str, Counter]): Keyword counts per category.
        labels (List[str]): List of category labels.
        remove_duplicates (bool): Whether to remove duplicates in clustered keywords.

    Returns:
        pd.DataFrame: Updated DataFrame with 'Clustered Keywords' and category columns.
    """
    logging.info("Generating clustered keywords and assigning categories...")
    # Build mapping from keywords to categories based on top keywords per category
    top_keywords_per_category = {label: set() for label in labels}
    # For the purpose of this example, we will assume some predefined keywords per category
    # In practice, you should define or compute these based on your data and requirements
    category_keywords = {
        'Type': {'shirt', 'dress', 'jeans', 'mask', 'boots', 'necklace'},
        'Color': {'black', 'blue', 'red', 'silver', 'gold', 'white'},
        'Wear Location': {'neck', 'foot', 'body', 'face'},
        'Material': {'rubber', 'metal', 'fabric', 'leather'},
        'Style': {'winter', 'venetian', 'western'}
    }
    # Update the top keywords per category
    for label in labels:
        top_keywords_per_category[label] = category_keywords.get(label, set())
    # Build a reverse mapping from keyword to categories
    keyword_to_categories = {}
    for label, keywords in top_keywords_per_category.items():
        for kw in keywords:
            if kw not in keyword_to_categories:
                keyword_to_categories[kw] = []
            keyword_to_categories[kw].append(label)
    # Initialize lists for categories
    category_keywords_list = {label: [] for label in labels}
    clustered_keywords_list = []
    for idx, row in per_row_df.iterrows():
        lemmatized_keywords = row['Lemmatized Keywords'].split(', ')
        clustered_keywords = [clustered_keyword_mapping.get(kw, kw) for kw in lemmatized_keywords]
        if remove_duplicates:
            # Remove duplicates while preserving order
            clustered_keywords = list(dict.fromkeys(clustered_keywords))
        clustered_keywords_list.append(', '.join(clustered_keywords))
        # Initialize assigned keywords for this row
        assigned_keywords = {label: 'Other' for label in labels}
        # Keep track of which categories have been assigned
        assigned_labels = set()
        # For each keyword
        for kw in clustered_keywords:
            categories = keyword_to_categories.get(kw, [])
            for label in categories:
                if label not in assigned_labels:
                    assigned_keywords[label] = kw
                    assigned_labels.add(label)
        # Add assigned keywords to category lists
        for label in labels:
            category_keywords_list[label].append(assigned_keywords[label])
    per_row_df['Clustered Keywords'] = clustered_keywords_list
    # Add the category columns to per_row_df
    for label in labels:
        per_row_df[label] = category_keywords_list[label]
    logging.info("Clustered keywords and categories generated successfully.")
    return per_row_df

def encode_categorical_features(input_df: pd.DataFrame, output_tsv_path: str, id_column: str = 'Image Name',
                                remove_extension: str = '.jpg', labels: List[str] = None, sep: str = '\t',
                                header: bool = False, index: bool = False):
    """
    Encodes categorical features using One-Hot Encoding and saves them to a TSV file.

    Parameters:
        input_df (pd.DataFrame): Input DataFrame containing the categorical features.
        output_tsv_path (str): Path to save the one-hot encoded TSV file.
        id_column (str): Name of the identifier column.
        remove_extension (str): Extension to remove from the identifier column values.
        labels (List[str]): List of columns to encode.
        sep (str): Delimiter to use in the TSV file.
        header (bool): Whether to include the header in the TSV file.
        index (bool): Whether to write row names (index).
    """
    logging.info("Starting one-hot encoding of categorical features...")
    if labels is None:
        raise ValueError("Labels for encoding must be provided.")
    # Read only the necessary columns
    columns_to_encode = labels
    columns_to_read = [id_column] + columns_to_encode
    df = input_df[columns_to_read]
    # Initialize OneHotEncoder
    enc = OneHotEncoder(handle_unknown='ignore', dtype=np.int8)
    # Fit and transform the categorical columns
    emb = enc.fit_transform(df[columns_to_encode])
    # Retrieve the names of the one-hot encoded features
    feature_names = enc.get_feature_names_out(columns_to_encode)
    # Create a DataFrame with the encoded features
    emb_df = pd.DataFrame(emb.toarray(), columns=feature_names, dtype=np.int8)
    # Concatenate the id_column with the one-hot encoded features
    result_df = pd.concat([df[[id_column]].reset_index(drop=True), emb_df.reset_index(drop=True)], axis=1)
    # Remove specified extension from id_column
    if remove_extension:
        result_df[id_column] = result_df[id_column].str.replace(remove_extension, '', regex=False)
    # Save to TSV file
    try:
        result_df.to_csv(output_tsv_path, index=index, header=header, sep=sep)
        logging.info(f"One-hot encoded features saved successfully to '{output_tsv_path}'.")
    except Exception as e:
        logging.error(f"Error saving one-hot encoded features: {e}")
        raise

if __name__ == "__main__":
    main()