STEP0_extract_entities.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "novel_path='./零之使魔.txt'\n",
    "output_dir='./output_zero/'\n",
    "extract_names_cache_path='./extract_names_cache.json'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(novel_path, 'r') as f:\n",
    "    all_lines = f.read().splitlines()\n",
    "all_lines = [line for line in all_lines if line.strip() != '']\n",
    "\n",
    "\n",
    "import os,json\n",
    "from tqdm.auto import tqdm\n",
    "extract_names_cache={'current_line':0,'char_freq':{}}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fc9576170dd54a2fb9e0cf42cf91e888",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading tf_model.h5:   0%|          | 0.00/1.11G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "All TF 2.0 model weights were used when initializing XLMRobertaForTokenClassification.\n",
      "\n",
      "All the weights of XLMRobertaForTokenClassification were initialized from the TF 2.0 model.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use XLMRobertaForTokenClassification for predictions without further training.\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForTokenClassification\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"opensource/extract_names\")\n",
    "model = AutoModelForTokenClassification.from_pretrained(\"opensource/extract_names\",from_tf=True)\n",
    "model=model.cuda().eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    她的手臂被抓住了\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'entity': 'ORG',\n",
       "  'score': 0.7558713,\n",
       "  'index': 1,\n",
       "  'word': '▁',\n",
       "  'start': 3,\n",
       "  'end': 4},\n",
       " {'entity': 'ORG',\n",
       "  'score': 0.76424843,\n",
       "  'index': 2,\n",
       "  'word': '她的',\n",
       "  'start': 4,\n",
       "  'end': 6},\n",
       " {'entity': 'ORG',\n",
       "  'score': 0.7591963,\n",
       "  'index': 3,\n",
       "  'word': '手臂',\n",
       "  'start': 6,\n",
       "  'end': 8},\n",
       " {'entity': 'ORG',\n",
       "  'score': 0.7379903,\n",
       "  'index': 4,\n",
       "  'word': '被',\n",
       "  'start': 8,\n",
       "  'end': 9},\n",
       " {'entity': 'ORG',\n",
       "  'score': 0.70303804,\n",
       "  'index': 5,\n",
       "  'word': '抓住',\n",
       "  'start': 9,\n",
       "  'end': 11},\n",
       " {'entity': 'ORG',\n",
       "  'score': 0.7389716,\n",
       "  'index': 6,\n",
       "  'word': '了',\n",
       "  'start': 11,\n",
       "  'end': 12}]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import pipeline\n",
    "classifier=pipeline('ner',model=model,tokenizer=tokenizer,device=0)\n",
    "text=all_lines[326]\n",
    "print(text)\n",
    "classifier(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tensorflow",
   "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.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}