MT_task.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "881ca06268d1761f",
   "metadata": {},
   "source": [
    "## Zero-shot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dfa01d5aaada82",
   "metadata": {},
   "outputs": [],
   "source": [
    "## ZeroShot learning for SIB-200 dataset\n",
    "\n",
    "from transformers import BloomForCausalLM\n",
    "from transformers import BloomTokenizerFast\n",
    "import os\n",
    "import pandas as pd\n",
    "import torch\n",
    "import warnings\n",
    "from transformers import XGLMTokenizer, XGLMForCausalLM\n",
    "\n",
    "model_address = \"bigscience/bloom-1b1\"\n",
    "\n",
    "# Filter out FutureWarning messages\n",
    "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
    "\n",
    "# Assuming data directory contains multiple subdirectories with test.tsv files\n",
    "data_directory = \"sib-200/data/annotated\"\n",
    "\n",
    "if model_address.startswith(\"facebook\"):\n",
    "    # Load ZeroShot learning model and tokenizer\n",
    "    model = XGLMForCausalLM.from_pretrained(model_address)\n",
    "    tokenizer = XGLMTokenizer.from_pretrained(model_address, padding_side='left')\n",
    "if model_address.startswith(\"bigscience\"):\n",
    "    # Load ZeroShot learning model and tokenizer\n",
    "    model = BloomForCausalLM.from_pretrained(model_address)\n",
    "    tokenizer = BloomTokenizerFast.from_pretrained(model_address)\n",
    "\n",
    "# Move model to GPU if available\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model.to(device)\n",
    "\n",
    "# Function to translate given texts to English\n",
    "def translate_batch(texts, input_language, output_language=\"English\", batch_size=5):\n",
    "    with torch.no_grad():  # Disable gradient calculation\n",
    "        generated_texts = []\n",
    "        for i in range(0, len(texts), batch_size):\n",
    "            batch_texts = texts[i:i+batch_size]\n",
    "            prompts = [f\"{input_language}: {text} \\n{output_language}:\" for text in batch_texts]\n",
    "            inputs = tokenizer(prompts, padding=True, return_tensors=\"pt\").to(device)\n",
    "            tokens = tokenizer.batch_encode_plus(prompts, add_special_tokens=True, padding=True, return_tensors=\"pt\")['input_ids']\n",
    "            result_length = tokens.shape[1] + 100\n",
    "            generated_batch = model.generate(inputs[\"input_ids\"], max_length=result_length)\n",
    "            for generated_text in generated_batch:\n",
    "                generated_texts.append(tokenizer.decode(generated_text, skip_special_tokens=True))\n",
    "            del inputs\n",
    "            torch.cuda.empty_cache()\n",
    "    return generated_texts\n",
    "\n",
    "# Output directory for saving DataFrames\n",
    "output_directory = \"MT-task/sib-200/zero-shot/\" + model_address[model_address.find('/')+1:] + \"/\"\n",
    "\n",
    "# Create the output directory if it doesn't exist\n",
    "os.makedirs(output_directory, exist_ok=True)\n",
    "\n",
    "language_df = pd.read_excel(\"SIB-200 languages - ACL.xlsx\")\n",
    "\n",
    "# Iterate through rows and compare predicted category with actual category\n",
    "for index, row in language_df.iterrows():\n",
    "    language = row['Language Name']\n",
    "    folder = row['Folder Name']\n",
    "    print(language)\n",
    "    if language == \"English\":\n",
    "        continue\n",
    "    if f\"{folder}.csv\" in os.listdir(output_directory):\n",
    "        continue\n",
    "    # Initialize an empty DataFrame to store results\n",
    "    results_df = pd.DataFrame(columns=['text', 'translated_text'])\n",
    "\n",
    "    subdir = os.path.join(data_directory, folder)\n",
    "    for file in os.listdir(subdir):\n",
    "        # Check if the file is a test.tsv file\n",
    "        if file.endswith(\"test.tsv\"):\n",
    "            file_path = os.path.join(subdir, file)\n",
    "            \n",
    "            # Check if the file is already present\n",
    "            results_file_path = os.path.join(output_directory, f'{folder}.csv')\n",
    "            if os.path.exists(results_file_path):\n",
    "                print(f\"Output file {results_file_path} already exists. Skipping...\")\n",
    "                continue\n",
    "\n",
    "            # Read the test.tsv file into a DataFrame\n",
    "            df = pd.read_csv(file_path, sep='\\t')\n",
    "\n",
    "            # Get all texts from the DataFrame\n",
    "            texts = df['text'].tolist()\n",
    "\n",
    "            # Predict translations using your ZeroShot learning model in batches\n",
    "            generated_texts = translate_batch(texts=texts, input_language=language)\n",
    "\n",
    "            # Append the results to the DataFrame\n",
    "            results_df['text'] = texts\n",
    "            results_df['translated_text'] = generated_texts\n",
    "\n",
    "            # Save the results DataFrame to a CSV file in the output directory\n",
    "            results_file_path = os.path.join(output_directory, f'{folder}.csv')\n",
    "            results_df.to_csv(results_file_path, index=False)\n",
    "\n",
    "            print(f\"Results saved to {results_file_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb76a7df8bc52e51",
   "metadata": {},
   "source": [
    "## Few-shot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf830aee3ac19c1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-21T04:15:08.784072Z",
     "start_time": "2024-04-21T04:05:30.783816Z"
    }
   },
   "outputs": [],
   "source": [
    "## FewShot learning for SIB-200 dataset\n",
    "\n",
    "from transformers import BloomForCausalLM\n",
    "from transformers import BloomTokenizerFast\n",
    "import os\n",
    "import pandas as pd\n",
    "import torch\n",
    "import warnings\n",
    "from transformers import XGLMTokenizer, XGLMForCausalLM\n",
    "\n",
    "model_address = \"bigscience/bloom-1b1\"\n",
    "n_shot = 2\n",
    "\n",
    "# Filter out FutureWarning messages\n",
    "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
    "\n",
    "# Assuming data directory contains multiple subdirectories with test.tsv files\n",
    "data_directory = \"sib-200/data/annotated\"\n",
    "\n",
    "if model_address.startswith(\"facebook\"):\n",
    "    # Load ZeroShot learning model and tokenizer\n",
    "    model = XGLMForCausalLM.from_pretrained(model_address)\n",
    "    tokenizer = XGLMTokenizer.from_pretrained(model_address, padding_side='left')\n",
    "if model_address.startswith(\"bigscience\"):\n",
    "    # Load ZeroShot learning model and tokenizer\n",
    "    model = BloomForCausalLM.from_pretrained(model_address)\n",
    "    tokenizer = BloomTokenizerFast.from_pretrained(model_address)\n",
    "\n",
    "# Move model to GPU if available\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model.to(device, dtype=torch.bfloat16)\n",
    "\n",
    "def few_shot_maker(input_folder, input_language, output_folder=\"eng_Latn\", output_language=\"English\", n_shots=2):\n",
    "    # Read the few-shot samples from the input folder\n",
    "    input_folder_path = os.path.join(data_directory, input_folder)\n",
    "    output_folder_path = os.path.join(data_directory, output_folder)\n",
    "    input_df = pd.read_csv(os.path.join(input_folder_path, \"train.tsv\"), sep='\\t')\n",
    "    output_df = pd.read_csv(os.path.join(output_folder_path, \"train.tsv\"), sep='\\t')\n",
    "    few_shot = \"\\n\\n\".join([f\"{input_language}: {input_df['text'][i]} \\n{output_language}: {output_df['text'][i]}\" for i in range(n_shots)])\n",
    "    return few_shot\n",
    "\n",
    "# Function to translate given texts to English\n",
    "def translate_batch(texts, few_shot_sample, input_language, output_language=\"English\", batch_size=16):\n",
    "    with torch.no_grad():  # Disable gradient calculation\n",
    "        generated_texts = []\n",
    "        for i in range(0, len(texts), batch_size):\n",
    "            batch_texts = texts[i:i+batch_size]\n",
    "            prompts = [f\"{few_shot_sample}\\n\\n{input_language}: {text} \\n{output_language}:\" for text in batch_texts]\n",
    "            inputs = tokenizer(prompts, padding=True, return_tensors=\"pt\").to(device)\n",
    "            tokens = tokenizer.batch_encode_plus(prompts, add_special_tokens=True, padding=True, return_tensors=\"pt\")['input_ids']\n",
    "            result_length = tokens.shape[1] + 100\n",
    "            generated_batch = model.generate(inputs[\"input_ids\"], max_length=result_length)\n",
    "            for generated_text in generated_batch:\n",
    "                generated_texts.append(tokenizer.decode(generated_text, skip_special_tokens=True))\n",
    "            del inputs\n",
    "            torch.cuda.empty_cache()\n",
    "    return generated_texts\n",
    "\n",
    "# Output directory for saving DataFrames\n",
    "output_directory = f\"MT-task/sib-200/{n_shot}-shot/\" + model_address[model_address.find('/')+1:] + \"/\"\n",
    "\n",
    "# Create the output directory if it doesn't exist\n",
    "os.makedirs(output_directory, exist_ok=True)\n",
    "\n",
    "language_df = pd.read_excel(\"SIB-200 languages - ACL.xlsx\")\n",
    "\n",
    "# Iterate through rows and compare predicted category with actual category\n",
    "for index, row in language_df.iterrows():\n",
    "    language = row['Language Name']\n",
    "    folder = row['Folder Name']\n",
    "    print(language)\n",
    "    if language == \"English\":\n",
    "        continue\n",
    "    if f\"{folder}.csv\" in os.listdir(output_directory):\n",
    "        continue\n",
    "    few_shot_sample = few_shot_maker(folder, language, n_shots=n_shot)    \n",
    "    \n",
    "    # Initialize an empty DataFrame to store results\n",
    "    results_df = pd.DataFrame(columns=['text', 'translated_text'])\n",
    "\n",
    "    subdir = os.path.join(data_directory, folder)\n",
    "    for file in os.listdir(subdir):\n",
    "        # Check if the file is a test.tsv file\n",
    "        if file.endswith(\"test.tsv\"):\n",
    "            file_path = os.path.join(subdir, file)\n",
    "\n",
    "            # Check if the file is already present\n",
    "            results_file_path = os.path.join(output_directory, f'{folder}.csv')\n",
    "            if os.path.exists(results_file_path):\n",
    "                print(f\"Output file {results_file_path} already exists. Skipping...\")\n",
    "                continue\n",
    "\n",
    "            # Read the test.tsv file into a DataFrame\n",
    "            df = pd.read_csv(file_path, sep='\\t')\n",
    "\n",
    "            # Get all texts from the DataFrame\n",
    "            texts = df['text'].tolist()\n",
    "\n",
    "            # Predict translations using your ZeroShot learning model in batches\n",
    "            generated_texts = translate_batch(texts=texts, few_shot_sample=few_shot_sample, input_language=language)\n",
    "\n",
    "            # Append the results to the DataFrame\n",
    "            results_df['text'] = texts\n",
    "            results_df['translated_text'] = generated_texts\n",
    "\n",
    "            # Save the results DataFrame to a CSV file in the output directory\n",
    "            results_file_path = os.path.join(output_directory, f'{folder}.csv')\n",
    "            results_df.to_csv(results_file_path, index=False)\n",
    "\n",
    "            print(f\"Results saved to {results_file_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae5e146c6abae6e3",
   "metadata": {},
   "source": [
    "## Remove duplicate sentences and clean the output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eebcbb76b7ddeb6b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-23T16:58:47.040929Z",
     "start_time": "2024-04-23T16:58:13.091638Z"
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os\n",
    "import nltk\n",
    "from nltk.tokenize import sent_tokenize\n",
    "\n",
    "# Function to remove duplicate sentences from a given text\n",
    "def remove_duplicate_sentences(text):\n",
    "    # Tokenize the text into sentences\n",
    "    sentences = sent_tokenize(text)\n",
    "\n",
    "    # Create a set to store unique sentences\n",
    "    unique_sentences = set()\n",
    "\n",
    "    # Iterate through each sentence\n",
    "    for sentence in sentences:\n",
    "        # Check if the sentence is not already in the set\n",
    "        if sentence not in unique_sentences:\n",
    "            # Add the unique sentence to the set\n",
    "            unique_sentences.add(sentence)\n",
    "\n",
    "    # Reconstruct the text with unique sentences\n",
    "    unique_text = ' '.join(unique_sentences)\n",
    "\n",
    "    return unique_text\n",
    "\n",
    "zero_shot_output = \"MT-task/sib-200/zero-shot/\"\n",
    "few_shot_output = \"MT-task/sib-200/2-shot/\"\n",
    "\n",
    "# Load all files in the zero-shot output directory and its subdirectories\n",
    "for root, dirs, files in os.walk(zero_shot_output):\n",
    "    for file in files:\n",
    "        if file.endswith(\".csv\"):\n",
    "            df = pd.read_csv(os.path.join(root, file))\n",
    "            # get text after \"/nEnglish:\" and remove repeated text\n",
    "            print(os.path.join(root, file))\n",
    "            df['clean_output'] = df['translated_text'].str.split(\"English:\").str[1].str.strip()\n",
    "            df['clean_output'] = df['clean_output'].apply(lambda x: remove_duplicate_sentences(x) if isinstance(x, str) else x)\n",
    "            # Save the cleaned DataFrame to a CSV file\n",
    "            df.to_csv(os.path.join(root, file), index=False)\n",
    "\n",
    "# # Load all files in the few-shot output directory and its subdirectories\n",
    "# for root, dirs, files in os.walk(few_shot_output):\n",
    "#     for file in files:\n",
    "#         if file.endswith(\".csv\"):\n",
    "#             df = pd.read_csv(os.path.join(root, file))\n",
    "#             # get text after \"/nEnglish:\" and remove repeated text\n",
    "#             print(os.path.join(root, file))\n",
    "#             df['clean_output'] = df['translated_text'].str.split(\"English:\").str[3].str.split(\"\\n\").str[0].str.strip()\n",
    "#             df['clean_output'] = df['clean_output'].apply(lambda x: remove_duplicate_sentences(x) if isinstance(x, str) else x)\n",
    "#             # Save the cleaned DataFrame to a CSV file\n",
    "#             df.to_csv(os.path.join(root, file), index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec4abe8aa3f644b0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}