transformers/GloVe.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "正在从http://d2l-data.s3-accelerate.amazonaws.com/glove.6B.50d.zip下载../data_sets/glove.6B.50d.zip...\n",
      "<__main__.TokenEmbedding object at 0x7fc2a4c6bbb0>\n",
      "3367 beautiful\n",
      "chips：cosine相似度=0.856\n",
      "intel：cosine相似度=0.749\n",
      "electronics：cosine相似度=0.749\n",
      "None\n",
      "babies：cosine相似度=0.839\n",
      "boy：cosine相似度=0.800\n",
      "girl：cosine相似度=0.792\n",
      "None\n",
      "lovely：cosine相似度=0.921\n",
      "gorgeous：cosine相似度=0.893\n",
      "wonderful：cosine相似度=0.830\n",
      "None\n",
      "daughter\n",
      "japan\n",
      "went\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import torch\n",
    "from torch import nn\n",
    "import sys\n",
    "sys.path.append('..')\n",
    "from tools import tool_pytorch_017 as ty\n",
    "\n",
    "#@save\n",
    "ty.DATA_HUB['glove.6b.50d'] = (ty.DATA_URL + 'glove.6B.50d.zip',\n",
    "                                '0b8703943ccdb6eb788e6f091b8946e82231bc4d')\n",
    "\n",
    "#@save\n",
    "ty.DATA_HUB['glove.6b.100d'] = (ty.DATA_URL + 'glove.6B.100d.zip',\n",
    "                                 'cd43bfb07e44e6f27cbcc7bc9ae3d80284fdaf5a')\n",
    "\n",
    "#@save\n",
    "ty.DATA_HUB['glove.42b.300d'] = (ty.DATA_URL + 'glove.42B.300d.zip',\n",
    "                                  'b5116e234e9eb9076672cfeabf5469f3eec904fa')\n",
    "\n",
    "#@save\n",
    "ty.DATA_HUB['wiki.en'] = (ty.DATA_URL + 'wiki.en.zip',\n",
    "                           'c1816da3821ae9f43899be655002f6c723e91b88')\n",
    "\n",
    "#@save\n",
    "class TokenEmbedding:\n",
    "    \"\"\"GloVe嵌入\"\"\"\n",
    "    def __init__(self, embedding_name):\n",
    "        self.idx_to_token, self.idx_to_vec = self._load_embedding(\n",
    "            embedding_name)\n",
    "        self.unknown_idx = 0\n",
    "        self.token_to_idx = {token: idx for idx, token in\n",
    "                             enumerate(self.idx_to_token)}\n",
    "\n",
    "    def _load_embedding(self, embedding_name):\n",
    "        idx_to_token, idx_to_vec = ['<unk>'], []\n",
    "        data_dir = ty.download_extract(embedding_name)\n",
    "        # GloVe网站：https://nlp.stanford.edu/projects/glove/\n",
    "        # fastText网站：https://fasttext.cc/\n",
    "        with open(os.path.join(data_dir, 'vec.txt'), 'r') as f:\n",
    "            for line in f:\n",
    "                elems = line.rstrip().split(' ')\n",
    "                token, elems = elems[0], [float(elem) for elem in elems[1:]]\n",
    "                # 跳过标题信息，例如fastText中的首行\n",
    "                if len(elems) > 1:\n",
    "                    idx_to_token.append(token)\n",
    "                    idx_to_vec.append(elems)\n",
    "        idx_to_vec = [[0] * len(idx_to_vec[0])] + idx_to_vec\n",
    "        return idx_to_token, torch.tensor(idx_to_vec)\n",
    "\n",
    "    def __getitem__(self, tokens):\n",
    "        indices = [self.token_to_idx.get(token, self.unknown_idx)\n",
    "                   for token in tokens]\n",
    "        vecs = self.idx_to_vec[torch.tensor(indices)]\n",
    "        return vecs\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.idx_to_token)\n",
    "\n",
    "# 加载50维GloVe嵌入（在维基百科的子集上预训练）。\n",
    "glove_6b50d = TokenEmbedding('glove.6b.50d')\n",
    "print(glove_6b50d)\n",
    "print(glove_6b50d.token_to_idx['beautiful'], glove_6b50d.idx_to_token[3367])\n",
    "\n",
    "def knn(W, x, k):\n",
    "    # 增加1e-9以获得数值稳定性\n",
    "    cos = torch.mv(W, x.reshape(-1,)) / (\n",
    "        torch.sqrt(torch.sum(W * W, axis=1) + 1e-9) *\n",
    "        torch.sqrt((x * x).sum()))\n",
    "    _, topk = torch.topk(cos, k=k)\n",
    "    return topk, [cos[int(i)] for i in topk]\n",
    "\n",
    "def get_similar_tokens(query_token, k, embed):\n",
    "    topk, cos = knn(embed.idx_to_vec, embed[[query_token]], k + 1)\n",
    "    for i, c in zip(topk[1:], cos[1:]):  # 排除输入词\n",
    "        print(f'{embed.idx_to_token[int(i)]}：cosine相似度={float(c):.3f}')\n",
    "        \n",
    "print(get_similar_tokens('chip', 3, glove_6b50d))\n",
    "print(get_similar_tokens('baby', 3, glove_6b50d))\n",
    "print(get_similar_tokens('beautiful', 3, glove_6b50d))\n",
    "\n",
    "def get_analogy(token_a, token_b, token_c, embed):\n",
    "    '''词类比'''\n",
    "    vecs = embed[[token_a, token_b, token_c]]\n",
    "    x = vecs[1] - vecs[0] + vecs[2]\n",
    "    topk, cos = knn(embed.idx_to_vec, x, 1)\n",
    "    return embed.idx_to_token[int(topk[0])]  # 删除未知词\n",
    "\n",
    "print(get_analogy('man', 'woman', 'son', glove_6b50d))\n",
    "print(get_analogy('beijing', 'china', 'tokyo', glove_6b50d))\n",
    "print(get_analogy('do', 'did', 'go', glove_6b50d))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3.8",
   "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.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}