Mandarin_A_maze_Alternative_Generation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/UMWordLab/multilingual_amaze/blob/main/Mandarin_A_maze_Alternative_Generation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Mandarin Maze Generation\n",
        "\n",
        "Lisa Levinson<sup>1</sup>, Yizhi Tang<sup>2</sup>, Lucy Yu-Chuan Chiang<sup>1</sup>, Wei-Jie Zhou<sup>1</sup>, Sohee Chung<sup>1</sup>\n",
        "\n",
        "(<sup>1</sup>University of Michigan, <sup>2</sup>Columbia University)\n",
        "\n",
        "**Summary.** We introduce an extension of the automatic maze task described in A-Maze (Boyce et al. 2020) for Mandarin.\n",
        "\n",
        "This script currently works for **simplified Mandarin** texts. We plan to share a variant for traditional Mandarin orthography in the near future."
      ],
      "metadata": {
        "id": "BXYd91y9TLlO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Script Setup\n",
        "Run the following cell and ignore any warnings. You must authorize Google Drive access for the script to work."
      ],
      "metadata": {
        "id": "pkdpE8VWYIAn"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4A9Ccsf4GDur"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "\n",
        "!pip install --upgrade -q gspread\n",
        "!python -m spacy download zh_core_web_md\n",
        "!pip install transformers\n",
        "!pip install spacy\n",
        "!pip install tdqm\n",
        "!pip install urllib3\n",
        "\n",
        "from tqdm.notebook import tqdm\n",
        "import spacy\n",
        "import torch\n",
        "import csv\n",
        "import random\n",
        "import gspread\n",
        "import pandas as pd\n",
        "import re\n",
        "import requests\n",
        "import urllib.request\n",
        "\n",
        "from spacy import displacy\n",
        "from collections import Counter\n",
        "from transformers import BertTokenizer, BertForMaskedLM\n",
        "from google.colab import drive\n",
        "from google.colab import auth\n",
        "from google.auth import default\n",
        "from google.colab import files\n",
        "\n",
        "tokenizer = BertTokenizer.from_pretrained(\"hfl/chinese-bert-wwm\")\n",
        "model = BertForMaskedLM.from_pretrained(\"hfl/chinese-bert-wwm\")\n",
        "nlp = spacy.load('zh_core_web_md')\n",
        "creds, _ = default()\n",
        "gc = gspread.authorize(creds)\n",
        "drive.mount('/content/drive')\n",
        "auth.authenticate_user()\n",
        "\n",
        "\n",
        "WINDOW_SIZE = 25"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## File Inputs and Variables\n",
        "\n",
        "As input, you will need a Google Sheet with a row for each stimulus sentence,  containing the words/phrases in separate columns (such as \"phrase1\", \"phrase2\", etc.). Each row should also contain an item label. If you want to use the matching tool to match words across sentences based on matching labels, items that you want to match must have the same item labels and you will need a column containing the matching labels. \n",
        "\n",
        "Once your stimuli are prepared, run the following cell and follow the prompts. After running, you will see a preview of how your input stimuli have been parsed. (See documentation for more detail)"
      ],
      "metadata": {
        "id": "3L6fDITtZAAs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "file_name = input(\"What is the url for the Google Sheet containing the stimuli?\\n\")\n",
        "sheet_name = input(\"What is the name of the sheet containing the stimuli (e.g., Sheet1)?\\n\")\n",
        "row_num = input(\"What row do your phrases start at? (Enter an integer - usually row 2 if you have a header row)\\n\")\n",
        "col_num = input(\"What column do your phrases start at? (Enter an integer)\\n\")\n",
        "col_end_num = input(\"What column do your phrases end at (Enter an integer)\\n\")\n",
        "lab_num = input(\"What column is your item labels in?\\n\")\n",
        "match_lab = input(\"What column are your matching labels in? (if none enter z)\\n\")\n",
        "outfile_name = input(\"What would you like your output file name to be?\\n\")\n",
        "\n",
        "\n",
        "row = int(row_num) - 1\n",
        "col = int(col_num) - 1\n",
        "lab = int(lab_num) - 1\n",
        "end = int(col_end_num) - 1\n",
        "\n",
        "worksheet = gc.open_by_url(file_name).worksheet(sheet_name)\n",
        "rows = worksheet.get_all_values()\n",
        "data = pd.DataFrame.from_records(rows)\n",
        "\n",
        "data"
      ],
      "metadata": {
        "id": "HQshj6EpGl6d"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Methods\n",
        "\n",
        "Run the following cell to define the necessary functions for generating alternatives and writing output."
      ],
      "metadata": {
        "id": "W--6J1OhZu8W"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def tokenization(sentence, split):\n",
        "  masked_list = []\n",
        "  inputs = tokenizer(sentence, add_special_tokens=True, return_tensors=\"pt\")\n",
        "  mask_index = 0 \n",
        "  encoding = inputs['input_ids'].clone()\n",
        "\n",
        "  for i in range(len(split)): \n",
        "    masked_list.append(inputs['input_ids'].clone())\n",
        "\n",
        "    for j in range(len(split[i])):\n",
        "      masked_list[i][0][mask_index + 1 + j] = 103\n",
        "\n",
        "    mask_index = mask_index + len(split[i])\n",
        "\n",
        "  return masked_list\n",
        "\n",
        "def find_list_of_words(word):\n",
        "  path = \"https://raw.githubusercontent.com/UMWordLab/multilingual_amaze/main/frequency/mandarin_simplified/\"\n",
        "\n",
        "  if len(word) == 1:\n",
        "    filename = \"one_character.csv\"\n",
        "\n",
        "  elif len(word) == 2:\n",
        "    filename = \"two_character.csv\"\n",
        "\n",
        "  elif len(word) == 3:\n",
        "    filename = \"three_character.csv\"\n",
        "\n",
        "  elif len(word) == 4:\n",
        "    filename = \"four_character.csv\"\n",
        "\n",
        "  else:\n",
        "    filename = \"long.csv\"\n",
        "\n",
        "  response = urllib.request.urlopen(path + filename)\n",
        "  lines = [l.decode('utf-8') for l in response.readlines()]\n",
        "  cr = csv.reader(lines)\n",
        "\n",
        "   #with open(filename, 'r') as csvfile:\n",
        "      #csvreader = csv.reader(csvfile)\n",
        "  rows = []\n",
        "  index = -1\n",
        "  i = 0\n",
        "\n",
        "  for row in cr:\n",
        "        rows.append(row)\n",
        "        if row[0] == word:\n",
        "          index = i\n",
        "        i = i + 1\n",
        "\n",
        "  if index == -1:\n",
        "        print(\"Warning: cannot find this word in our frequency list\")\n",
        "        res = []\n",
        "        random.sample(rows, WINDOW_SIZE * 2)\n",
        "\n",
        "        for item in random.sample(rows, WINDOW_SIZE * 2):\n",
        "          for token in nlp(item[0]):\n",
        "            res.append((item[0], item[1], token.pos_))\n",
        "\n",
        "        return res\n",
        "\n",
        "  similar_freq = []\n",
        "  for j in range(max(0, index - WINDOW_SIZE), min(len(rows), index + WINDOW_SIZE)):\n",
        "        for token in nlp(rows[j][0]):\n",
        "          similar_freq.append((rows[j][0], rows[j][1], token.pos_))\n",
        "\n",
        "  return similar_freq\n",
        "\n",
        "def calculate_surprisal(sentence, word, token, start_position, verbose_mode):\n",
        "    inputs = tokenizer(sentence, add_special_tokens=True, return_tensors=\"pt\")\n",
        "    inputs['input_ids'] = token  \n",
        "    outputs = model(**inputs) \n",
        "    similar = find_list_of_words(word)\n",
        "    surprisal_list = []\n",
        "\n",
        "    for triple in similar:\n",
        "      i = 0\n",
        "      prob = 0\n",
        "      for character in triple[0]:\n",
        "        embeddings = tokenizer.convert_tokens_to_ids(character)\n",
        "        actual_position = start_position + i + 1\n",
        "        word_weights = outputs[0][0][actual_position].squeeze().div(1.0).exp()\n",
        "\n",
        "        if i == 0:\n",
        "          prob = (word_weights / sum(word_weights))[embeddings]\n",
        "\n",
        "        else:\n",
        "          prob = prob * (word_weights / sum(word_weights))[embeddings]\n",
        "\n",
        "        i = i + 1\n",
        "\n",
        "      if verbose_mode:\n",
        "        print(triple[0], prob)\n",
        "      surprisal_list.append(-1 * torch.log2(prob))\n",
        "    \n",
        "    max_val = max(surprisal_list)\n",
        "    max_index = surprisal_list.index(max_val)\n",
        "\n",
        "    while True:\n",
        "      count = 0\n",
        "      flag = False\n",
        "\n",
        "      for token in nlp(word):\n",
        "        if similar[max_index][2] == token.pos_:\n",
        "          count += 1\n",
        "          surprisal_list.pop(max_index)\n",
        "          similar.pop(max_index)\n",
        "          max_val = max(surprisal_list)\n",
        "          max_index = surprisal_list.index(max_val)\n",
        "\n",
        "        else:\n",
        "          flag = True\n",
        "      \n",
        "      if flag or len(surprisal_list) == 0 or count > 10:\n",
        "        break\n",
        "\n",
        "    return similar[max_index]\n",
        "\n",
        "def find_alternative(sentence, split):\n",
        "  token_list = tokenization(sentence, split)\n",
        "  start_position = len(split[0])\n",
        "  result = []\n",
        "\n",
        "  for i in range(1, len(split)):\n",
        "    alter = calculate_surprisal(sentence, split[i], token_list[i], start_position, verbose_mode=False)\n",
        "    result.append(alter)\n",
        "    start_position = start_position + len(split[i])\n",
        "\n",
        "  return result\n",
        "\n",
        "def df_into_list(row_start, input_data, col_start, col_end):\n",
        "  list_of_splits = []\n",
        "  list_of_sentences = []\n",
        "  list_of_labels = []\n",
        "\n",
        "  for i in range(row_start,len(input_data)):\n",
        "    lst = []\n",
        "\n",
        "    for j in range(col_start,col_end):\n",
        "      if data[j][i] != '':\n",
        "        lst.append(data[j][i])\n",
        "\n",
        "    list_of_splits.append(lst)\n",
        "\n",
        "  for i in range(0, len(list_of_splits)):\n",
        "    list_of_sentences.append(''.join(list_of_splits[i]))\n",
        "    list_of_labels.append(data[0][i + 1])\n",
        "\n",
        "  print(len(list_of_sentences))\n",
        "\n",
        "  return list_of_sentences, list_of_splits, list_of_labels\n",
        "\n",
        "def write_to_csv(sentences_list, splits_list, labels_list, output_file):\n",
        "  f = open(output_file, 'a')\n",
        "  writer = csv.writer(f)\n",
        "  alternative_splits_list = []\n",
        "\n",
        "  for i in tqdm(range(len(sentences_list))):\n",
        "    result = find_alternative(sentences_list[i], splits_list[i])\n",
        "    alt_splits = []\n",
        "    alt_splits.append(splits_list[i][0])\n",
        "\n",
        "    for alternatives in result:\n",
        "        alt_splits.append(alternatives[0])\n",
        "\n",
        "    splits_list[i].insert(0, labels_list[i])\n",
        "    alt_splits.insert(0, labels_list[i] + '.alt')\n",
        "\n",
        "    alternative_splits_list.append(alt_splits)\n",
        "    writer.writerow(splits_list[i])\n",
        "    writer.writerow(alt_splits)\n",
        "\n",
        "    print(\"finished sentence \" + str(i + 1) + \" out of \" + str(len(sentences_list)))\n",
        "\n",
        "  f.close()\n",
        "\n",
        "def matcher(alternatives_splits_list, label_sentence_list, labels_used_list = ['alt', 'verb']): \n",
        "  distractors = {}\n",
        "  distractor_supplier_sent = alternatives_splits_list[0]\n",
        "\n",
        "  for i in range(len(label_sentence_list[0].split())):\n",
        "    if label_sentence_list[0].split()[i] in labels_used_list:\n",
        "\n",
        "      distractors[label_sentence_list[0].split()[i]] = distractor_supplier_sent[i]\n",
        "  \n",
        "  for i in range(len(alternatives_splits_list)):\n",
        "    words = alternatives_splits_list[i]\n",
        "\n",
        "    for word_pos in range(1, len(words)):\n",
        "\n",
        "      if label_sentence_list[i].split()[word_pos] in distractors:\n",
        "\n",
        "        alternatives_splits_list[i][word_pos] = distractors[label_sentence_list[i].split()[word_pos]]\n",
        "  \n",
        "  print(distractors)\n",
        "  return(alternatives_splits_list)\n",
        "\n",
        "def find_longest_list(lists):\n",
        "    longest_list = []\n",
        "    max_length = 0\n",
        "    \n",
        "    for lst in lists:\n",
        "        length = len(lst)\n",
        "        \n",
        "        if length > max_length:\n",
        "            longest_list = lst\n",
        "            max_length = length\n",
        "    \n",
        "    return max_length\n",
        "\n",
        "def find_largest_number(string):\n",
        "    numbers = []\n",
        "    current_number = \"\"\n",
        "\n",
        "    for char in string:\n",
        "        if char.isdigit():\n",
        "            current_number += char\n",
        "        else:\n",
        "            if current_number != \"\":\n",
        "                numbers.append(int(current_number))\n",
        "                current_number = \"\"\n",
        "\n",
        "    if current_number != \"\":\n",
        "        numbers.append(int(current_number))\n",
        "\n",
        "    if len(numbers) == 0:\n",
        "        return None\n",
        "\n",
        "    return max(numbers)\n",
        "\n",
        "def get_unique_numbers(df, column):\n",
        "    unique_numbers = set()\n",
        "    \n",
        "    for value in df[column]:\n",
        "        if isinstance(value, (int, float)):\n",
        "            unique_numbers.add(int(value))\n",
        "        elif isinstance(value, str):\n",
        "            numbers = re.findall(r'\\d+', value)\n",
        "            for number in numbers:\n",
        "                unique_numbers.add(int(number))\n",
        "\n",
        "    return list(unique_numbers)\n",
        "\n",
        "def find_rows_with_integer(df, column, integer):\n",
        "    matching_rows = []\n",
        "\n",
        "    for index, row in df.iterrows():\n",
        "        if str(integer) in str(row[column]):\n",
        "            matching_rows.append(row)\n",
        "\n",
        "    return pd.DataFrame(matching_rows)\n",
        "\n",
        "def flatten_list_of_list_of_lists(list_of_list_of_lists):\n",
        "    flattened_list = [item for sublist in list_of_list_of_lists for item in sublist]\n",
        "    return flattened_list\n",
        "\n",
        "def select_rows(df, match_strings, output_column):\n",
        "    selected_rows = df.loc[df.iloc[:, lab].isin(match_strings), output_column].tolist()\n",
        "    return selected_rows\n",
        "\n",
        "def write_list_of_lists_to_csv(list_of_lists, filename):\n",
        "    with open(filename, 'w', newline='') as file:\n",
        "        writer = csv.writer(file)\n",
        "\n",
        "        for row in list_of_lists:\n",
        "            writer.writerow(row)\n",
        "\n",
        "def find_common_words(sentences):\n",
        "    if not sentences:\n",
        "        return []\n",
        "\n",
        "    word_sets = [set(sentence.lower().split()) for sentence in sentences]\n",
        "    common_words = set.intersection(*word_sets)\n",
        "    print(\"Labels being matched are: \" + str(common_words))\n",
        "\n",
        "    for word_set in word_sets[1:]:\n",
        "        common_words &= word_set\n",
        "    return sorted(list(common_words))\n",
        "\n",
        "def match_labels_exist(num):\n",
        "  if num == \"z\":\n",
        "    return\n",
        "  else:\n",
        "    return int(num)"
      ],
      "metadata": {
        "id": "-SHpo1FtGmuB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the following cell to make sure it is the correct number of sentences."
      ],
      "metadata": {
        "id": "NpBwpJ8tQZwg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sentence_list, split_list, labeled_list = df_into_list(row, data, col, end)"
      ],
      "metadata": {
        "id": "AXAwxLdBJZ_t"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Generate Alternatives"
      ],
      "metadata": {
        "id": "SWarGhxoDjG6"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "This section is where the alternatives are actually generated, which is the computationally intensive portion. **This step will take approximately 4 minutes PER SENTENCE. Please plan accordingly.** If you have a large set of sentences, you may want to generate alternatives in batches to make sure that the server does not timeout or run out of memory before completing the job. \n",
        "\n",
        "Run the following cell to write the alternatives to the output file specified in the **File Inputs and Variables** step. The file can be found if you click the folder icon on the left edge of the Colab window, which will open the \"Files\" tab. You can download the file by highlighting it and clicking the three dot menu on the right side. **Files are deleted from the Colab Notebook when the session disconnects/times out - make sure to save the output to your computer or a cloud drive.**\n",
        "\n"
      ],
      "metadata": {
        "id": "nOC8eej3Qhri"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "write_to_csv(sentence_list, split_list, labeled_list, outfile_name)"
      ],
      "metadata": {
        "id": "GQPHXVegaHY6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Within-Item Alternative Matching (Optional)\n",
        "You should now see a file in your google colab's files tab named \"your_inputted_output_name\".txt. We now use these generated alternatives to match the alternatives based on inputted matching labels and item labels."
      ],
      "metadata": {
        "id": "WMh6t695j50h"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the following cell and input a new filename different from your original output that you want the ouptut to be named. The new output with matched alternatives will also be saved to the files tab. "
      ],
      "metadata": {
        "id": "i-yOqx2v5LLi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "distractor_wks = pd.read_csv('outs.txt', header = None)\n",
        "matched_alts_outfile = input(\"What do you want the outfile to be named? \\n\")\n",
        "match_lab_col = match_labels_exist(match_lab) - 1\n",
        "distractor_list = list(data[match_lab_col])"
      ],
      "metadata": {
        "id": "sj0RKDCYQl1n"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the following two cells"
      ],
      "metadata": {
        "id": "Eopk3FYa5Z5a"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "unique_labels = get_unique_numbers(distractor_wks, 0)\n",
        "alternatives_groups = []\n",
        "\n",
        "for i in unique_labels:\n",
        "  alternatives_groups.append(find_rows_with_integer(distractor_wks, 0, i))\n",
        "\n",
        "print(unique_labels)\n",
        "print(alternatives_groups)"
      ],
      "metadata": {
        "id": "Xyv7DgVg9Hf1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "alts_list = []\n",
        "alts_label_list = []\n",
        "\n",
        "for i in range(len(alternatives_groups)):\n",
        "  distractor_wks_temp = alternatives_groups[i][alternatives_groups[i][0].str.contains('alt')]\n",
        "  alts_labels = distractor_wks_temp.iloc[:, 0].tolist()\n",
        "  new_list = [row[1:].tolist() for index, row in distractor_wks_temp.iterrows()]\n",
        "  alts_labels = [word[:-4] for word in alts_labels]\n",
        "\n",
        "  matching_label_list = find_common_words(select_rows(data, alts_labels, match_lab_col))\n",
        "  alts_list_group = matcher(new_list, select_rows(data, alts_labels, match_lab_col), matching_label_list)\n",
        "  alts_list.append(alts_list_group)\n",
        "  alts_label_list.append(alts_labels)\n",
        "\n",
        "alts_list = flatten_list_of_list_of_lists(alts_list)\n",
        "alts_label_list = flatten_list_of_list_of_lists(alts_label_list)\n",
        "\n",
        "for i in range(len(alts_label_list)):\n",
        "  alts_list[i].insert(0, alts_label_list[i] + \".matchedalt\")"
      ],
      "metadata": {
        "id": "71ZGkuDM1L9l"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "result = []\n",
        "\n",
        "for i in range(len(alts_list)):\n",
        "  split_list[i].insert(0, labeled_list[i])\n",
        "  result.append(split_list[i])\n",
        "  result.append(alts_list[i])\n",
        "\n",
        "write_list_of_lists_to_csv(result, matched_alts_outfile)"
      ],
      "metadata": {
        "id": "1pfipN3WF7U1"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}