get_openai_embeddings.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d39ce6d1-3ab0-4edf-9568-1472238c3d2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from openai import OpenAI\n",
    "from pie_data import get_dataset\n",
    "#import tiktoken\n",
    "import numpy as np\n",
    "import os\n",
    "import pickle\n",
    "import yaml\n",
    "import json\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfc0f052-d490-495b-a92a-45eafe2fe490",
   "metadata": {},
   "source": [
    "### (!) Specifiy your API Key\n",
    "\n",
    "Create a file `api_key.txt` with your OpenAI API Key to run this notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "128c368d-7b69-4d44-b23d-9e356b1d7eb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('api_key.txt', 'r') as file:\n",
    "    api_key = file.read().rstrip()\n",
    "\n",
    "client = OpenAI(api_key=api_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6422a88b-5f29-44fd-bb60-a016ca51a401",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_model = \"text-embedding-3-small\"\n",
    "embedding_encoding = \"cl100k_base\"\n",
    "max_tokens = 800  # the maximum for text-embedding-3-small is 8191"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "791452c1-2403-4ad9-a4d4-232968924fc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_raw_embeddings(texts, save_file, embedding_model):\n",
    "    if os.path.exists(save_file):\n",
    "        with open(save_file, 'rb') as handle:\n",
    "            out_chunks = pickle.load(handle)\n",
    "    else:\n",
    "        out_chunks = []\n",
    "        for i in tqdm(range(0, len(texts), 10)):\n",
    "            out_chunks.append(client.embeddings.create(input=texts[i:i+10], model=embedding_model))\n",
    "\n",
    "    return out_chunks\n",
    "\n",
    "\n",
    "def get_embeddings(texts, save_file_raw, save_file_np, embedding_model):\n",
    "    if os.path.exists(save_file_np):\n",
    "        with open(save_file_np, 'rb') as handle:\n",
    "            emb_dict = pickle.load(handle)\n",
    "            assert emb_dict['model'] == embedding_model\n",
    "            embeddings = emb_dict['embeddings']\n",
    "            assert len(embeddings) == len(texts), (\"found %i embeddings for %i texts\" % (len(embeddings), len(texts)))\n",
    "    else:\n",
    "        out_chunks = get_raw_embeddings(texts, save_file_raw, embedding_model)\n",
    "\n",
    "        with open(save_file_raw, 'wb') as handle:\n",
    "            pickle.dump(out_chunks, handle)\n",
    "            \n",
    "        # remove chunks from raw embedding list while filling new list (necessary for larger datasets)\n",
    "        embeddings = []\n",
    "        while len(out_chunks) > 0:\n",
    "            chunk = out_chunks.pop(0)\n",
    "            for elem in chunk.data:\n",
    "                embeddings.append(elem.embedding)\n",
    "        \n",
    "        assert len(embeddings) == len(texts), (\"found %i embeddings for %i texts\" % (len(embeddings), len(texts)))\n",
    "\n",
    "        emb_arr = np.asarray(embeddings)\n",
    "        saved = {'model': embedding_model, 'embeddings': emb_arr}\n",
    "        \n",
    "        with open(save_file_np, 'wb') as handle:\n",
    "            pickle.dump(saved, handle)\n",
    "        \n",
    "    return embeddings\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2fb2d3f-8bf1-4699-b968-694ca58b72ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_split_emb(data, split, dataset_name, embedding_model):\n",
    "    print(\"got %i samples for split %s\" % (len(data), split))\n",
    "    save_file_raw = ('embeddings/%s_%s_%s_raw_output.pickle' % (dataset_name, split, embedding_model))\n",
    "    save_file_np = ('embeddings/%s_%s_%s.pickle' % (dataset_name, split, embedding_model))\n",
    "    embeddings = get_embeddings(data, save_file_raw, save_file_np, embedding_model)\n",
    "    return embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98accde0-38f8-406c-96f1-fc11d10104af",
   "metadata": {},
   "source": [
    "### Embed training and test splits of the datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d6f20a8-883d-4f00-aad4-094101484955",
   "metadata": {},
   "outputs": [],
   "source": [
    "# bios supervised\n",
    "dataset_name = 'bios-supervised'\n",
    "bios_dir = '../../data/bios_huggingface_merge.pkl'\n",
    "X_train, y_train, X_test, y_test, n_classes, multi_label, class_weights, protected_attr_dict = get_dataset(dataset_name, local_dir=bios_dir)\n",
    "\n",
    "emb = {}\n",
    "emb['train'] = get_split_emb(X_train, 'train', dataset_name, embedding_model)\n",
    "emb['test'] = get_split_emb(X_test, 'test', dataset_name, embedding_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4f892a3-6e83-4470-b499-6540396b0d27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# twitterAAE\n",
    "dataset_name = 'twitterAAE'\n",
    "X_train, y_train, X_test, y_test, n_classes, multi_label, class_weights, protected_attr_dict = get_dataset(dataset_name)\n",
    "\n",
    "emb = {}\n",
    "emb['test'] = get_split_emb(X_test, 'test', dataset_name, embedding_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f4c455c-a04a-494d-8190-3c282e22d7ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# crowspairs\n",
    "dataset_name = 'crows_pairs'\n",
    "X_train, y_train, X_test, y_test, n_classes, multi_label, class_weights, protected_attr_dict = get_dataset(dataset_name)\n",
    "\n",
    "emb = {}\n",
    "emb['test'] = get_split_emb(X_test, 'test', dataset_name, embedding_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c04de773-f79a-467c-a917-026ea260ab3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# bios unsupervised\n",
    "dataset_name = 'bios-unsupervised'\n",
    "X_train, y_train, X_test, y_test, n_classes, multi_label, class_weights, protected_attr_dict = get_dataset(dataset_name)\n",
    "\n",
    "emb = {}\n",
    "emb['train'] = get_split_emb(X_train, 'train', dataset_name, embedding_model)\n",
    "emb['test'] = get_split_emb(X_test, 'test', dataset_name, embedding_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86086f7f-07d1-4e5a-a1ec-6b2a2a2a2a24",
   "metadata": {},
   "outputs": [],
   "source": [
    "# jigsaw\n",
    "dataset_name = 'jigsaw'\n",
    "local_dir = '../../data/jigsaw_bias'\n",
    "X_train, y_train, X_test, y_test, n_classes, multi_label, class_weights, protected_attr_dict = get_dataset(dataset_name, local_dir)\n",
    "\n",
    "emb = {}\n",
    "emb['train'] = get_split_emb(X_train, 'train', dataset_name, embedding_model)\n",
    "emb['test'] = get_split_emb(X_test, 'test', dataset_name, embedding_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af9699dc-82c3-426f-a9ca-1ef9f2cd8c1e",
   "metadata": {},
   "source": [
    "### Create dictionary with words/phrases used as defining terms in the experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e271b53-1104-467e-abd3-b551ce6cb748",
   "metadata": {},
   "outputs": [],
   "source": [
    "dict_empty = 'embeddings/word_phrase_dict_empty.pickle'\n",
    "\n",
    "with open(dict_empty, 'rb') as handle:\n",
    "    word_phrase_emb_dict_empty = pickle.load(handle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0c631e2-1e25-4208-88e1-3c3ff938b3d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# update word phrase dictionary given an experiment config\n",
    "CONFIG_FILE = 'experiments/configs/esann25/experiment_config.json'\n",
    "\n",
    "with open(CONFIG_FILE, 'r') as f:\n",
    "    config = json.load(f)\n",
    "\n",
    "eval_setup_file = config['bias_space_eval_config']\n",
    "with open(eval_setup_file, 'r') as stream:\n",
    "    eval_setups_by_attr = yaml.safe_load(stream)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1413236a-966b-46e0-a7d0-ac5e01544c5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(word_phrase_emb_dict_empty)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f677f524-ad08-4597-b9d1-4288eba0ae19",
   "metadata": {},
   "outputs": [],
   "source": [
    "for attr, content in eval_setups_by_attr.items():\n",
    "    assert len(content['defining_terms']) == 1\n",
    "\n",
    "    for group_terms in content['defining_terms'][0]:\n",
    "        for term in group_terms:\n",
    "            if not term in word_phrase_emb_dict_empty.keys():\n",
    "                word_phrase_emb_dict_empty[term] = None\n",
    "\n",
    "with open(dict_empty, 'wb') as handle:\n",
    "     pickle.dump(word_phrase_emb_dict_empty, handle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8b2485f-2691-434d-a3e5-19c7a9606b85",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(word_phrase_emb_dict_empty)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e53b014e-ee31-404d-98f8-3b66100facdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load or create new word/phrase dictionary for the current embedding model\n",
    "dict_emb = ('embeddings/word_phrase_dict_%s.pickle' % embedding_model)\n",
    "\n",
    "# load current state of dictionary (if available)\n",
    "if os.path.exists(dict_emb):\n",
    "    with open(dict_emb, 'rb') as handle:\n",
    "        loaded_dict = pickle.load(handle)\n",
    "        prev_model = loaded_dict['model']\n",
    "        assert prev_model == embedding_model\n",
    "        word_phrase_emb_dict = loaded_dict['emb_dict']\n",
    "else:\n",
    "    word_phrase_emb_dict = {}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17e0298d-c7b3-4aa4-a73e-b356a22e00ed",
   "metadata": {},
   "source": [
    "### Embed the terms and phrases from the dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea0c7ca7-6d98-4330-9865-d20614c84511",
   "metadata": {},
   "outputs": [],
   "source": [
    "# query word/phrase embedding for current embedding model\n",
    "save_dict = {'model': embedding_model, 'emb_dict': word_phrase_emb_dict}\n",
    "\n",
    "for term, emb in word_phrase_emb_dict_empty.items():\n",
    "    if term in word_phrase_emb_dict.keys() and word_phrase_emb_dict[term] is not None:\n",
    "        # embedding for this term or phrase already exists\n",
    "        continue\n",
    "    else:\n",
    "        # call api\n",
    "        print(\"call api for %s\" % term)\n",
    "        emb = client.embeddings.create(input=[term], model=embedding_model).data[0].embedding\n",
    "        save_dict['emb_dict'][term] = emb\n",
    "\n",
    "with open(dict_emb, 'wb') as handle:\n",
    "    pickle.dump(save_dict, handle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15342862-349e-4840-bc23-df4281e72a4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(save_dict['emb_dict'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed7b50d1-927b-4896-99fa-1371afe5745b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48ae6b07-90d6-4730-b1f3-f13f3c6da81e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}