PyTorch_NLP_NER.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4fd512ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "import torch\n",
    "from torch.utils.data import DataLoader,Dataset\n",
    "import torch.nn as nn\n",
    "device = torch.device(\"mps\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7818d413",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('/Users/aadityajoshi/Downloads/ner.csv',index_col=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ab469aaa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Sentence</th>\n",
       "      <th>POS</th>\n",
       "      <th>Tag</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Thousands of demonstrators have marched throug...</td>\n",
       "      <td>['NNS', 'IN', 'NNS', 'VBP', 'VBN', 'IN', 'NNP'...</td>\n",
       "      <td>['O', 'O', 'O', 'O', 'O', 'O', 'B-geo', 'O', '...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Families of soldiers killed in the conflict jo...</td>\n",
       "      <td>['NNS', 'IN', 'NNS', 'VBN', 'IN', 'DT', 'NN', ...</td>\n",
       "      <td>['O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>They marched from the Houses of Parliament to ...</td>\n",
       "      <td>['PRP', 'VBD', 'IN', 'DT', 'NNS', 'IN', 'NN', ...</td>\n",
       "      <td>['O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Police put the number of marchers at 10,000 wh...</td>\n",
       "      <td>['NNS', 'VBD', 'DT', 'NN', 'IN', 'NNS', 'IN', ...</td>\n",
       "      <td>['O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>The protest comes on the eve of the annual con...</td>\n",
       "      <td>['DT', 'NN', 'VBZ', 'IN', 'DT', 'NN', 'IN', 'D...</td>\n",
       "      <td>['O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            Sentence  \\\n",
       "0  Thousands of demonstrators have marched throug...   \n",
       "1  Families of soldiers killed in the conflict jo...   \n",
       "2  They marched from the Houses of Parliament to ...   \n",
       "3  Police put the number of marchers at 10,000 wh...   \n",
       "4  The protest comes on the eve of the annual con...   \n",
       "\n",
       "                                                 POS  \\\n",
       "0  ['NNS', 'IN', 'NNS', 'VBP', 'VBN', 'IN', 'NNP'...   \n",
       "1  ['NNS', 'IN', 'NNS', 'VBN', 'IN', 'DT', 'NN', ...   \n",
       "2  ['PRP', 'VBD', 'IN', 'DT', 'NNS', 'IN', 'NN', ...   \n",
       "3  ['NNS', 'VBD', 'DT', 'NN', 'IN', 'NNS', 'IN', ...   \n",
       "4  ['DT', 'NN', 'VBZ', 'IN', 'DT', 'NN', 'IN', 'D...   \n",
       "\n",
       "                                                 Tag  \n",
       "0  ['O', 'O', 'O', 'O', 'O', 'O', 'B-geo', 'O', '...  \n",
       "1  ['O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', ...  \n",
       "2  ['O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', ...  \n",
       "3  ['O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', ...  \n",
       "4  ['O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', ...  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.index = range(len(df))\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f159afcf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 47959 entries, 0 to 47958\n",
      "Data columns (total 3 columns):\n",
      " #   Column    Non-Null Count  Dtype \n",
      "---  ------    --------------  ----- \n",
      " 0   Sentence  47959 non-null  object\n",
      " 1   POS       47959 non-null  object\n",
      " 2   Tag       47959 non-null  object\n",
      "dtypes: object(3)\n",
      "memory usage: 1.1+ MB\n"
     ]
    }
   ],
   "source": [
    "df.info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b5cce0c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df.sample(frac=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "560bd0c4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The government announced it would eliminate 5,00,000 state jobs by March 2011 and has expanded opportunities for self-employment .'"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df['text'].iloc[893]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4b579e64",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.dropna(inplace=True)\n",
    "df.drop_duplicates(inplace=True)\n",
    "blanks = []\n",
    "for i,s,p,t in df.itertuples():\n",
    "    if type(s) == str:\n",
    "        if s.isspace()==True:\n",
    "            blanks.append(i)\n",
    "blanks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ff8c44ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "df['text'] = df['Sentence']\n",
    "df.drop('Sentence',axis=1,inplace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31a325ae",
   "metadata": {},
   "source": [
    "### Create lists of everything"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "247f9b0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "texts = df['text'].apply(lambda x : x.split()).tolist()\n",
    "pos_tags = df['POS'].apply(eval).tolist()\n",
    "ner_tags = df['Tag'].apply(eval).tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e183c995",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4796"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(ner_tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e11e308b",
   "metadata": {},
   "source": [
    "### create vocabs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "39b4e092",
   "metadata": {},
   "outputs": [],
   "source": [
    "word_vocab = {'<PAD>':0,'<UNK>':1}\n",
    "pos_vocab = {'<PAD>':0,'<UNK>':1}\n",
    "tag_vocab = {'<PAD>':0,'<UNK>':1}\n",
    "for sent in texts:\n",
    "    for word in sent:\n",
    "        if word not in word_vocab:\n",
    "            word_vocab[word] = len(word_vocab)\n",
    "for pos_seq in pos_tags:\n",
    "    for pos in pos_seq:\n",
    "        if pos not in pos_vocab:\n",
    "            pos_vocab[pos] = len(pos_vocab)\n",
    "for tag_seq in ner_tags:\n",
    "    for ner in tag_seq:\n",
    "        if ner not in tag_vocab:\n",
    "            tag_vocab[ner] = len(tag_vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ce76806e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "19"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(tag_vocab)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8660f892",
   "metadata": {},
   "source": [
    "### using vocab convert into indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7c240f60",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_ind,pos_ind,tag_ind = [],[],[]\n",
    "for i in range(len(texts)):\n",
    "    sent_ind = [word_vocab.get(text,word_vocab['<UNK>']) for text in texts[i]]\n",
    "    text_ind.append(sent_ind)\n",
    "    pos_tag_idx = [pos_vocab.get(pos,pos_vocab['<UNK>']) for pos in pos_tags[i]]\n",
    "    pos_ind.append(pos_tag_idx)\n",
    "    ner_tag_idx = [tag_vocab.get(tag,tag_vocab['<UNK>']) for tag in ner_tags[i]]\n",
    "    tag_ind.append(ner_tag_idx)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a48a4ed5",
   "metadata": {},
   "source": [
    "### padding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "eb25de6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad_sequence(sequences):\n",
    "    max_len = max([len(seq) for seq in sequences])\n",
    "    padded_sequences = []\n",
    "    for seq in sequences:  \n",
    "        padded_seq = seq + [0] * (max_len - len(seq))\n",
    "        padded_sequences.append(padded_seq)\n",
    "    return torch.tensor(padded_sequences)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0cb65c14",
   "metadata": {},
   "outputs": [],
   "source": [
    "padded_text = pad_sequence(text_ind)\n",
    "padded_pos = pad_sequence(pos_ind)\n",
    "padded_ner = pad_sequence(tag_ind)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d2c28c5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "lengths = torch.tensor([len(seq) for seq in text_ind])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "78f3769e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4796"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(lengths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "0895ac5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "processed_data = {\n",
    "    'texts':padded_text,\n",
    "    'pos':padded_pos,\n",
    "    'tags':padded_ner,\n",
    "    'lengths':lengths,\n",
    "    'vocabularies':{\n",
    "        'word_vocab' : word_vocab,\n",
    "        'pos_vocab' : pos_vocab,\n",
    "        'tag_vocab': tag_vocab\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62f85185",
   "metadata": {},
   "source": [
    "### Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "68175e25",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NERDataset(Dataset):\n",
    "    def __init__(self,processed_data):\n",
    "        self.texts = processed_data['texts']\n",
    "        self.pos = processed_data['pos']\n",
    "        self.tags = processed_data['tags']\n",
    "        self.lengths = processed_data['lengths']\n",
    "    def __len__(self):\n",
    "        return len(self.texts)\n",
    "    def __getitem__(self,idx):\n",
    "        return {'text':self.texts[idx],\n",
    "               'pos':self.pos[idx],\n",
    "               'tag':self.tags[idx],\n",
    "               'length':self.lengths[idx]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "c3f86e02",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = NERDataset(processed_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "05c9f989",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.NERDataset at 0x175842f50>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "1a32c387",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4796"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "391e6320",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_size = len(dataset)\n",
    "train_size = int(0.8 * total_size)  # 80% for training\n",
    "val_size = total_size - train_size  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "bcfc3629",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import random_split\n",
    "train_dataset, test_dataset = random_split(\n",
    "    dataset, \n",
    "    [train_size, val_size],\n",
    "    generator=torch.Generator().manual_seed(42)  # For reproducibility\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "d30a47ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate_fn(batch):\n",
    "    # Sort batch by sequence length\n",
    "    batch = sorted(batch, key=lambda x: x['length'], reverse=True)\n",
    "    \n",
    "    # Stack all tensors\n",
    "    texts = torch.stack([item['text'] for item in batch])\n",
    "    pos = torch.stack([item['pos'] for item in batch])\n",
    "    tags = torch.stack([item['tag'] for item in batch])\n",
    "    lengths = torch.tensor([item['length'] for item in batch])\n",
    "    \n",
    "    return {\n",
    "        'text': texts,\n",
    "        'pos': pos,\n",
    "        'tag': tags,\n",
    "        'length': lengths\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "a59575be",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = DataLoader(train_dataset,batch_size=128,shuffle=True,collate_fn=collate_fn)\n",
    "test_loader = DataLoader(test_dataset,batch_size=128,shuffle=False,collate_fn=collate_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "435b79e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTM(nn.Module):\n",
    "    def __init__(self,vocab_size,pos_size,num_tags,embedding_dim=100,pos_emb_dim=20,hidden_size=128,num_layers=1,dropout=0.4):\n",
    "        super(LSTM,self).__init__()\n",
    "        self.embeddings = nn.Embedding(vocab_size,embedding_dim)\n",
    "        self.pos_embeddings = nn.Embedding(pos_size,pos_emb_dim)\n",
    "        total_dim = embedding_dim + pos_emb_dim\n",
    "        self.lstm = nn.LSTM(total_dim,hidden_size,num_layers=num_layers,bidirectional=True,batch_first=True)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.fc = nn.Linear(2*hidden_size,num_tags)\n",
    "    def forward(self, text, pos, lengths):\n",
    "        word_embeds = self.embeddings(text)\n",
    "        pos_embeds = self.pos_embeddings(pos)\n",
    "        embeds = torch.cat([word_embeds, pos_embeds], dim=2)\n",
    "        packed_embeds = pack_padded_sequence(\n",
    "            embeds, lengths.cpu(), batch_first=True, enforce_sorted=False\n",
    "        )\n",
    "        packed_output, _ = self.lstm(packed_embeds)\n",
    "        output, _ = pad_packed_sequence(packed_output, batch_first=True, \n",
    "                                      total_length=text.size(1))  # Pad to original length\n",
    "        tag_scores = self.fc(output)\n",
    "        return tag_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "3274ab6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = LSTM(\n",
    "    vocab_size=len(word_vocab),\n",
    "    pos_size=len(pos_vocab),\n",
    "    num_tags=len(tag_vocab)\n",
    ").to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "f9593728",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LSTM(\n",
       "  (embeddings): Embedding(12013, 100)\n",
       "  (pos_embeddings): Embedding(43, 20)\n",
       "  (lstm): LSTM(120, 128, batch_first=True, bidirectional=True)\n",
       "  (dropout): Dropout(p=0.4, inplace=False)\n",
       "  (fc): Linear(in_features=256, out_features=19, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "910b8751",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 10\n",
    "optimizer = torch.optim.Adam(model.parameters(),lr=0.001)\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "ab416ed3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1/10, Loss: 0.0066\n",
      "Epoch: 2/10, Loss: 0.0055\n",
      "Epoch: 3/10, Loss: 0.0047\n",
      "Epoch: 4/10, Loss: 0.0041\n",
      "Epoch: 5/10, Loss: 0.0035\n",
      "Epoch: 6/10, Loss: 0.0031\n",
      "Epoch: 7/10, Loss: 0.0027\n",
      "Epoch: 8/10, Loss: 0.0024\n",
      "Epoch: 9/10, Loss: 0.0021\n",
      "Epoch: 10/10, Loss: 0.0019\n"
     ]
    }
   ],
   "source": [
    "from torch.nn.utils.rnn import pack_padded_sequence,pad_packed_sequence\n",
    "model.train()\n",
    "for epoch in range(num_epochs):\n",
    "    total_loss = 0\n",
    "    \n",
    "    for batch in train_loader:\n",
    "        # Get batch data\n",
    "        text = batch['text'].to(device)\n",
    "        pos = batch['pos'].to(device)\n",
    "        tags = batch['tag'].to(device)\n",
    "        lengths = batch['length']\n",
    "        \n",
    "        # Forward pass\n",
    "        tag_scores = model(text, pos, lengths)\n",
    "        \n",
    "        # Reshape predictions and targets\n",
    "        tag_scores = tag_scores.view(-1, tag_scores.size(-1))  # (batch_size * seq_len, num_tags)\n",
    "        tags = tags.view(-1)  # (batch_size * seq_len)\n",
    "        \n",
    "        # Calculate loss\n",
    "        loss = criterion(tag_scores, tags)\n",
    "        \n",
    "        # Backward pass and optimization\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        total_loss += loss.item()\n",
    "    \n",
    "    # Print epoch statistics\n",
    "    avg_loss = total_loss / len(train_loader)\n",
    "    print(f\"Epoch: {epoch+1}/{num_epochs}, Loss: {avg_loss:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "a56b2fed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Accuracy: 0.9466\n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "correct = 0\n",
    "total = 0\n",
    "device = torch.device(\"mps\")\n",
    "with torch.no_grad():\n",
    "    for batch in test_loader:\n",
    "        # Move batch data to device\n",
    "        text = batch['text'].to(device)\n",
    "        pos = batch['pos'].to(device)\n",
    "        tags = batch['tag'].to(device)\n",
    "        lengths = batch['length']  # lengths stay on CPU\n",
    "        \n",
    "        # Get model predictions\n",
    "        outputs = model(text, pos, lengths)\n",
    "        predictions = torch.argmax(outputs, dim=2)\n",
    "        \n",
    "        # Calculate accuracy (ignoring padding)\n",
    "        mask = tags != 0\n",
    "        correct += (predictions[mask] == tags[mask]).sum().item()\n",
    "        total += mask.sum().item()\n",
    "\n",
    "# Calculate and print accuracy\n",
    "accuracy = correct / total\n",
    "print(f\"Test Accuracy: {accuracy:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "2ec91334",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Harry: O\n",
      "potter: O\n",
      "is: O\n",
      "main: O\n",
      "character: O\n",
      "of: O\n",
      "the: O\n",
      "novel: O\n"
     ]
    }
   ],
   "source": [
    "def predict_ner(model, text, word_vocab, pos_vocab, tag_vocab, nlp, device='mps'):\n",
    "\n",
    "    # Process the text\n",
    "    doc = nlp(text)\n",
    "    words = [token.text for token in doc]\n",
    "    pos_tags = [token.pos_ for token in doc]\n",
    "    \n",
    "    # Convert to indices\n",
    "    word_indices = [word_vocab.get(word, word_vocab['<UNK>']) for word in words]\n",
    "    pos_indices = [pos_vocab.get(pos, pos_vocab['<UNK>']) for pos in pos_tags]\n",
    "    \n",
    "    # Convert to tensors and add batch dimension\n",
    "    text_tensor = torch.tensor([word_indices], device=device)\n",
    "    pos_tensor = torch.tensor([pos_indices], device=device)\n",
    "    lengths = torch.tensor([len(words)])\n",
    "    \n",
    "    # Get predictions\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        outputs = model(text_tensor, pos_tensor, lengths)\n",
    "        predictions = torch.argmax(outputs, dim=2)[0]\n",
    "    \n",
    "    # Convert indices back to tags\n",
    "    rev_tag_vocab = {v: k for k, v in tag_vocab.items()}\n",
    "    predicted_tags = [rev_tag_vocab[idx.item()] for idx in predictions]\n",
    "    \n",
    "    # Return word-tag pairs\n",
    "    return list(zip(words, predicted_tags))\n",
    "\n",
    "# Example usage:\n",
    "\n",
    "import spacy\n",
    "nlp = spacy.load('en_core_web_sm')\n",
    "\n",
    "# Test a new sentence\n",
    "test_sentence = \"Harry potter is main character of the novel\"\n",
    "results = predict_ner(\n",
    "    model, \n",
    "    test_sentence,\n",
    "    processed_data['vocabularies']['word_vocab'],\n",
    "    processed_data['vocabularies']['pos_vocab'],\n",
    "    processed_data['vocabularies']['tag_vocab'],\n",
    "    nlp,\n",
    "    device\n",
    ")\n",
    "\n",
    "# Print results\n",
    "for word, tag in results:\n",
    "    print(f\"{word}: {tag}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f88c538",
   "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.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}