Add address extraction functionality using BERT model

.gitignore

@@ -0,0 +1,2 @@
+venv
+condaenv

README.md

@@ -1,3 +1,98 @@
 ---
 license: mit
+base_model: dbmdz/bert-base-turkish-cased
+tags:
+- ner
+- token-classification
+- pytorch
+- turkish
+- tr
+- dbmdz
+- bert
+- bert-base-cased
+- bert-base-turkish-cased
+widget:
+- text: "Bağlarbaşı Mahallesi, Zübeyde Hanım Caddesi No: 10 / 3 34710 Üsküdar/İstanbul"
 ---
+
+# address-extraction
+
+![Next Geography](https://nextgeography.com/wp-content/uploads/2022/02/next-geo-logo-1.png)
+
+This is a simple library to extract addresses from text. The train.py file contains the code to train but is just included for reference, not to be run. The model is trained on our own dataset of addresses, which is not included in this repo. There is also predict.py which is a simple script to run the model on a single address.
+
+The model is based on [dbmdz/bert-base-turkish-cased](https://huggingface.co/dbmdz/bert-base-turkish-cased) from [Hugging Face](https://huggingface.co/).
+
+## Example Results
+
+```
+(g:\projects\address-extraction\venv) G:\projects\address-extraction>python predict.py
+Osmangazi Mahallesi, Hoca Ahmet Yesevi Cd. No:34, 16050 Osmangazi/Bursa
+Osmangazi                                   Mahalle 98.65%
+Hoca Ahmet Yesevi                             Cadde 97.63%
+34                                    Bina Numarası 98.92%
+16050                                    Posta Kodu 97.83%
+Osmangazi                                      İlçe 98.97%
+Bursa                                            İl 99.21%
+Average Score:  0.9902257982053255
+Labels Found:  6
+----------------------------------------------------------------------
+Karşıyaka Mahallesi, Mavişehir Caddesi No: 91, Daire 4, 35540 Karşıyaka/İzmir
+Karşıyaka                                   Mahalle 99.11%
+Mavişehir                                     Cadde 97.16%
+91                                    Bina Numarası 98.73%
+4                                               Kat 29.06%
+35540                                    Posta Kodu 98.65%
+Karşıyaka                                      İlçe 99.17%
+İzmir                                            İl 99.16%
+Average Score:  0.9237866433043229
+Labels Found:  7
+----------------------------------------------------------------------
+Selçuklu Mahallesi, Atatürk Bulvarı No: 55, 42050 Selçuklu/Konya
+Selçuklu                                    Mahalle 98.67%
+Atatürk                                       Cadde 57.06%
+55                                    Bina Numarası 98.94%
+42050                                    Posta Kodu 98.81%
+Selçuklu                                       İlçe 99.06%
+Konya                                            İl 99.22%
+Average Score:  0.9659512996673584
+Labels Found:  6
+----------------------------------------------------------------------
+Alsancak Mahallesi, 1475. Sk. No:3, 35220 Konak/İzmir
+Alsancak                                    Mahalle 99.38%
+1475                                          Sokak 96.04%
+3                                     Bina Numarası 98.06%
+35220                                    Posta Kodu 98.75%
+Konak                                          İlçe 99.23%
+İzmir                                            İl 99.16%
+Average Score:  0.9909308176291617
+Labels Found:  6
+----------------------------------------------------------------------
+Kocatepe Mahallesi, Yaşam Caddesi 3. Sokak No:4, 06420 Bayrampaşa/İstanbul
+Kocatepe                                    Mahalle 99.46%
+Yaşam                                         Cadde 94.07%
+3                                             Sokak 84.07%
+4                                     Bina Numarası 98.42%
+06420                                    Posta Kodu 98.54%
+Bayrampaşa                                     İlçe 98.97%
+İstanbul                                         İl 98.98%
+Average Score:  0.9832726591511777
+Labels Found:  7
+----------------------------------------------------------------------
+```
+
+## Installation & Usage
+
+The environment.yml file contains the conda environment used to run the model. Environment is configured to use cuda enabled gpus but should work with no gpus too. To run the model, you can use the following commands:
+
+```bash
+conda env create -f environment.yml -p ./condaenv
+conda activate ./condaenv
+
+python predict.py
+```
+
+
+## License
+
+This project is licensed under the terms of the MIT license.

environment.yml

@@ -0,0 +1,148 @@
+name: address-extraction
+channels:
+  - pytorch
+  - nvidia
+  - conda-forge
+  - defaults
+dependencies:
+  - blas=2.120
+  - blas-devel=3.9.0
+  - brotli-python=1.1.0
+  - bzip2=1.0.8
+  - ca-certificates=2023.11.17
+  - certifi=2023.11.17
+  - charset-normalizer=3.3.2
+  - cuda-cccl=12.3.101
+  - cuda-cudart=12.1.105
+  - cuda-cudart-dev=12.1.105
+  - cuda-cupti=12.1.105
+  - cuda-libraries=12.1.0
+  - cuda-libraries-dev=12.1.0
+  - cuda-nvrtc=12.1.105
+  - cuda-nvrtc-dev=12.1.105
+  - cuda-nvtx=12.1.105
+  - cuda-opencl=12.3.101
+  - cuda-opencl-dev=12.3.101
+  - cuda-profiler-api=12.3.101
+  - cuda-runtime=12.1.0
+  - filelock=3.13.1
+  - freetype=2.12.1
+  - idna=3.6
+  - intel-openmp=2023.2.0
+  - jinja2=3.1.3
+  - lcms2=2.16
+  - lerc=4.0.0
+  - libblas=3.9.0
+  - libcblas=3.9.0
+  - libcublas=12.1.0.26
+  - libcublas-dev=12.1.0.26
+  - libcufft=11.0.2.4
+  - libcufft-dev=11.0.2.4
+  - libcurand=10.3.4.107
+  - libcurand-dev=10.3.4.107
+  - libcusolver=11.4.4.55
+  - libcusolver-dev=11.4.4.55
+  - libcusparse=12.0.2.55
+  - libcusparse-dev=12.0.2.55
+  - libdeflate=1.19
+  - libffi=3.4.2
+  - libhwloc=2.9.3
+  - libiconv=1.17
+  - libjpeg-turbo=3.0.0
+  - liblapack=3.9.0
+  - liblapacke=3.9.0
+  - libnpp=12.0.2.50
+  - libnpp-dev=12.0.2.50
+  - libnvjitlink=12.1.105
+  - libnvjitlink-dev=12.1.105
+  - libnvjpeg=12.1.1.14
+  - libnvjpeg-dev=12.1.1.14
+  - libpng=1.6.39
+  - libsqlite=3.44.2
+  - libtiff=4.6.0
+  - libuv=1.44.2
+  - libwebp-base=1.3.2
+  - libxcb=1.15
+  - libxml2=2.12.4
+  - libzlib=1.2.13
+  - m2w64-gcc-libgfortran=5.3.0
+  - m2w64-gcc-libs=5.3.0
+  - m2w64-gcc-libs-core=5.3.0
+  - m2w64-gmp=6.1.0
+  - m2w64-libwinpthread-git=5.0.0.4634.697f757
+  - markupsafe=2.1.4
+  - mkl=2023.2.0
+  - mkl-devel=2023.2.0
+  - mkl-include=2023.2.0
+  - mpmath=1.3.0
+  - msys2-conda-epoch=20160418
+  - networkx=3.2.1
+  - numpy=1.26.3
+  - openjpeg=2.5.0
+  - openssl=3.2.0
+  - pillow=10.2.0
+  - pip=23.3.2
+  - pthread-stubs=0.4
+  - pthreads-win32=2.9.1
+  - pysocks=1.7.1
+  - python=3.10.13
+  - python_abi=3.10
+  - pytorch=2.1.2
+  - pytorch-cuda=12.1
+  - pytorch-mutex=1.0
+  - pyyaml=6.0.1
+  - requests=2.31.0
+  - setuptools=69.0.3
+  - sympy=1.12
+  - tbb=2021.11.0
+  - tk=8.6.13
+  - typing_extensions=4.9.0
+  - ucrt=10.0.22621.0
+  - urllib3=2.1.0
+  - vc=14.3
+  - vc14_runtime=14.38.33130
+  - vs2015_runtime=14.38.33130
+  - wheel=0.42.0
+  - win_inet_pton=1.1.0
+  - xorg-libxau=1.0.11
+  - xorg-libxdmcp=1.1.3
+  - xz=5.2.6
+  - yaml=0.2.5
+  - zstd=1.5.5
+  - pip:
+      - accelerate==0.26.1
+      - aiohttp==3.9.1
+      - aiosignal==1.3.1
+      - async-timeout==4.0.3
+      - attrs==23.2.0
+      - colorama==0.4.6
+      - datasets==2.16.1
+      - dill==0.3.7
+      - frozenlist==1.4.1
+      - fsspec==2023.10.0
+      - huggingface-hub==0.20.3
+      - joblib==1.3.2
+      - multidict==6.0.4
+      - multiprocess==0.70.15
+      - orjson==3.9.12
+      - packaging==23.2
+      - pandas==2.2.0
+      - psutil==5.9.8
+      - pyarrow==15.0.0
+      - pyarrow-hotfix==0.6
+      - python-dateutil==2.8.2
+      - pytz==2023.3.post1
+      - regex==2023.12.25
+      - safetensors==0.4.1
+      - scikit-learn==1.4.0
+      - scipy==1.12.0
+      - six==1.16.0
+      - threadpoolctl==3.2.0
+      - tokenizers==0.15.1
+      - torchaudio==2.1.2
+      - torchvision==0.16.2
+      - tqdm==4.66.1
+      - transformers==4.37.0
+      - tzdata==2023.4
+      - xxhash==3.4.1
+      - yarl==1.9.4

labels.json

@@ -0,0 +1 @@
+{"1": "B-\u00dclke", "2": "I-\u00dclke", "3": "B-\u0130l", "4": "I-\u0130l", "5": "B-\u0130l\u00e7e", "6": "I-\u0130l\u00e7e", "7": "B-Mahalle", "8": "I-Mahalle", "9": "B-Cadde", "10": "I-Cadde", "11": "B-Sokak", "12": "I-Sokak", "13": "B-Bina Ad\u0131", "14": "I-Bina Ad\u0131", "15": "B-Bina Numaras\u0131", "16": "I-Bina Numaras\u0131", "17": "B-Yer Ad\u0131", "18": "I-Yer Ad\u0131", "19": "B-Site", "20": "I-Site", "21": "B-Adres Detay", "22": "I-Adres Detay", "23": "B-Blok No", "24": "I-Blok No", "25": "B-Bulvar", "26": "I-Bulvar", "27": "B-Daire No", "28": "I-Daire No", "29": "B-Posta Kodu", "30": "I-Posta Kodu", "31": "B-Kat", "32": "I-Kat", "0": "O"}

model/config.json

@@ -0,0 +1,95 @@
+{
+  "_name_or_path": "dbmdz/bert-base-turkish-cased",
+  "architectures": [
+    "BertForTokenClassification"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "classifier_dropout": null,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1",
+    "2": "LABEL_2",
+    "3": "LABEL_3",
+    "4": "LABEL_4",
+    "5": "LABEL_5",
+    "6": "LABEL_6",
+    "7": "LABEL_7",
+    "8": "LABEL_8",
+    "9": "LABEL_9",
+    "10": "LABEL_10",
+    "11": "LABEL_11",
+    "12": "LABEL_12",
+    "13": "LABEL_13",
+    "14": "LABEL_14",
+    "15": "LABEL_15",
+    "16": "LABEL_16",
+    "17": "LABEL_17",
+    "18": "LABEL_18",
+    "19": "LABEL_19",
+    "20": "LABEL_20",
+    "21": "LABEL_21",
+    "22": "LABEL_22",
+    "23": "LABEL_23",
+    "24": "LABEL_24",
+    "25": "LABEL_25",
+    "26": "LABEL_26",
+    "27": "LABEL_27",
+    "28": "LABEL_28",
+    "29": "LABEL_29",
+    "30": "LABEL_30",
+    "31": "LABEL_31",
+    "32": "LABEL_32"
+  },
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1,
+    "LABEL_10": 10,
+    "LABEL_11": 11,
+    "LABEL_12": 12,
+    "LABEL_13": 13,
+    "LABEL_14": 14,
+    "LABEL_15": 15,
+    "LABEL_16": 16,
+    "LABEL_17": 17,
+    "LABEL_18": 18,
+    "LABEL_19": 19,
+    "LABEL_2": 2,
+    "LABEL_20": 20,
+    "LABEL_21": 21,
+    "LABEL_22": 22,
+    "LABEL_23": 23,
+    "LABEL_24": 24,
+    "LABEL_25": 25,
+    "LABEL_26": 26,
+    "LABEL_27": 27,
+    "LABEL_28": 28,
+    "LABEL_29": 29,
+    "LABEL_3": 3,
+    "LABEL_30": 30,
+    "LABEL_31": 31,
+    "LABEL_32": 32,
+    "LABEL_4": 4,
+    "LABEL_5": 5,
+    "LABEL_6": 6,
+    "LABEL_7": 7,
+    "LABEL_8": 8,
+    "LABEL_9": 9
+  },
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.37.0",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 32000
+}

model/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e9ad463f26bde0807766033059922cb64c633ca3cedd5637d98607a5a906eaeb
+size 440231868

model/special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

model/tokenizer_config.json

@@ -0,0 +1,58 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "do_basic_tokenize": true,
+  "do_lower_case": false,
+  "mask_token": "[MASK]",
+  "max_len": 512,
+  "model_max_length": 512,
+  "never_split": null,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]"
+}

model/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52ec3649ec97e031b39c9b67ba30182a38c953393e011aa2fefb33da76b0ad9c
+size 4664

predict.py

@@ -0,0 +1,81 @@
+import torch
+import orjson
+from transformers import pipeline
+from transformers import BertTokenizerFast, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("dbmdz/bert-base-turkish-cased")
+
+with open("labels.json", "r") as f:
+    id_to_label = {int(k): v for k, v in orjson.loads(f.read()).items()}
+
+nlp = pipeline(
+    "ner",
+    model="./model",
+    tokenizer=tokenizer,
+    device=0 if torch.cuda.is_available() else -1,
+)
+
+def get_entities(tokens):
+    entities = []
+    entity = None
+    for token in tokens:
+        label_id = int(token["entity"][6:])
+        label = id_to_label[label_id]
+        if label.startswith("B-"):
+            if entity:
+                entity["score"] /= entity["token_count"]
+                entities.append(entity)
+            entity = {
+                "label": label[2:],
+                "ranges": [token["start"], token["end"]],
+                "score": token["score"],
+                "token_count": 1,
+            }
+        elif label.startswith("I-"):
+            if entity and entity["label"] == label[2:]:
+                entity["ranges"][1] = token["end"]
+                entity["token_count"] += 1
+                entity["score"] += token["score"]
+            else:
+                if entity:
+                    entity["ranges"][1] = token["end"]
+                    entity["token_count"] += 1
+                    entity["score"] += token["score"]
+                    entity["score"] /= entity["token_count"]
+                    entities.append(entity)
+                entity = None
+        else:
+            if entity:
+                entity["score"] /= entity["token_count"]
+                entities.append(entity)
+            entity = None
+    if entity:
+        entity["score"] /= entity["token_count"]
+        entities.append(entity)
+    return entities
+
+def process(text):
+    nlp_output = nlp(text)
+    entities = get_entities(nlp_output)
+    for entity in entities:
+        print(f"{text[entity['ranges'][0]:entity['ranges'][1]]:<35} {entity['label']:>15} {entity['score'] * 100:.2f}%")
+    print("Average Score: ", sum([token["score"] for token in nlp_output]) / len(nlp_output))
+    print("Labels Found: ", len(entities))
+    print("-" * 70)
+
+if __name__ == "__main__":
+    examples = [
+        "Osmangazi Mahallesi, Hoca Ahmet Yesevi Cd. No:34, 16050 Osmangazi/Bursa",
+        "Karşıyaka Mahallesi, Mavişehir Caddesi No: 91, Daire 4, 35540 Karşıyaka/İzmir",
+        "Selçuklu Mahallesi, Atatürk Bulvarı No: 55, 42050 Selçuklu/Konya",
+        "Alsancak Mahallesi, 1475. Sk. No:3, 35220 Konak/İzmir",
+        "Kocatepe Mahallesi, Yaşam Caddesi 3. Sokak No:4, 06420 Bayrampaşa/İstanbul",
+    ]
+    for example in examples:
+        print(example)
+        process(example)
+    while True:
+        text = input("Enter text: ")
+        if not text:
+            break
+        process(text)

train.py

@@ -0,0 +1,199 @@
+import io
+import requests
+import json
+import time
+import torch
+import orjson
+import zipfile
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from transformers import BertTokenizerFast, BertForTokenClassification, Trainer, TrainingArguments, BertConfig
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+
+API_URL = "http://dockerbase.duo:8000"
+PROJECT_ID = 1
+
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+def load_data():
+    
+    res = requests.post(
+        API_URL + "/v1/auth/login/",
+        json={"username": "admin", "password": "123"}
+    )
+    token = res.json()["key"]
+
+    res = requests.post(API_URL + "/v1/projects/1/download",
+        json={"format":"JSONL","exportApproved": True},
+        headers={"Authorization": "Token " + token}
+    )
+    task_id = res.json()["task_id"]
+
+
+    ready = False
+    print("Waiting for export task to be ready.", end="")
+    while not ready:
+        res = requests.get(
+            API_URL + "/v1/tasks/status/" + str(task_id),
+            headers={"Authorization": "Token " + token}
+        )
+        ready = res.json()["ready"]
+        if not ready:
+            time.sleep(1)
+            print(".", end="")
+    print("")
+
+    res = requests.get(
+        API_URL + f"/v1/projects/{PROJECT_ID}/download",
+        params={"taskId": task_id},
+        headers={"Authorization": "Token " + token}
+    )
+
+    zip_file = io.BytesIO(res.content)
+    with zipfile.ZipFile(zip_file, "r") as zip_ref:
+        data = zip_ref.read("admin.jsonl").decode("utf-8")
+    
+    res = requests.get(
+        API_URL + f"/v1/projects/{PROJECT_ID}/span-types",
+        headers={"Authorization": "Token " + token}
+    )
+
+    labels = res.json()
+
+    return labels, [orjson.loads(line) for line in data.split("\n") if line]
+
+labels, data = load_data()
+label_to_id = {}
+for i, label in enumerate(labels):
+    label_to_id["B-" + label["text"]] = i * 2 + 1
+    label_to_id["I-" + label["text"]] = i * 2 + 2
+label_to_id["O"] = 0
+id_to_label = {v: k for k, v in label_to_id.items()}
+
+tokenizer = AutoTokenizer.from_pretrained("dbmdz/bert-base-turkish-cased")
+model = AutoModelForTokenClassification.from_pretrained("dbmdz/bert-base-turkish-cased", num_labels=len(label_to_id)).to(device)
+
+
+from datasets import DatasetDict, Dataset
+
+
+def preprocess_data(item, tokenizer, label_to_id):
+    text = item['text']
+
+    inputs = tokenizer(
+        text,
+        return_offsets_mapping=True,
+        return_tensors="pt",
+        truncation=True,
+        padding='max_length',
+        max_length=128,
+    )
+
+    input_ids = inputs["input_ids"]
+    attention_mask = inputs["attention_mask"]
+    offset_mapping = inputs["offset_mapping"]
+
+    labels = ["O"] * 128
+    last_label = "O"
+    for token_idx, [off_start, off_end] in enumerate(offset_mapping[0]):
+        if off_start == off_end:
+            continue
+
+        for start, end, label in item['label']:
+            if start <= off_start and off_end <= end:
+                if last_label == label:
+                    labels[token_idx] = "I-" + label
+                else:
+                    labels[token_idx] = "B-" + label
+                last_label = label
+                break
+    
+    # Convert labels to ids
+    labels = [label_to_id[label] for label in labels]
+
+    return {
+        "input_ids": input_ids.flatten(),
+        "attention_mask": attention_mask.flatten(),
+        "labels": labels,
+    }
+
+
+class AddressDataset(Dataset):
+    def __init__(self, dataset):
+        self.dataset = dataset
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, index):
+        item = self.dataset[index]
+        return {key: torch.tensor(val) for key, val in item.items()}
+
+
+
+dataset = Dataset.from_generator(
+    lambda: (preprocess_data(item, tokenizer, label_to_id) for item in data),
+)
+
+dataset = dataset.train_test_split(test_size=0.2)
+dataset = DatasetDict({
+    "train": dataset["train"],
+    "test": dataset["test"]
+})
+
+
+training_args = TrainingArguments(
+    output_dir="./results",
+    num_train_epochs=35,
+    per_device_train_batch_size=32,
+    per_device_eval_batch_size=32,
+    # logging_dir="./logs",
+    # logging_first_step=True,
+    # evaluation_strategy="epoch",
+    # save_strategy="epoch",
+    logging_strategy="epoch",
+    # load_best_model_at_end=True,
+)
+
+from sklearn.preprocessing import MultiLabelBinarizer
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
+
+def compute_metrics(pred, id_to_label):
+    labels = pred.label_ids
+    preds = pred.predictions.argmax(-1)
+
+    labels = [[id_to_label[label_id] for label_id in label_ids] for label_ids in labels]
+    preds = [[id_to_label[pred] for pred in preds] for preds in preds]
+
+    labels = [label for label in labels if label != "O"]
+    preds = [pred for pred in preds if pred != "O"]
+
+    mlb = MultiLabelBinarizer()
+    mlb.fit([id_to_label.values()])
+    labels = mlb.transform(labels)
+    preds = mlb.transform(preds)
+
+    return {
+        "accuracy": accuracy_score(labels, preds),
+        "precision": precision_score(labels, preds, average="micro"),
+        "recall": recall_score(labels, preds, average="micro"),
+        "f1": f1_score(labels, preds, average="micro"),
+    }
+
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=dataset["train"],
+    eval_dataset=dataset["test"],
+    tokenizer=tokenizer,
+    compute_metrics=lambda p: compute_metrics(p, id_to_label),
+)
+
+trainer.train()
+trainer.evaluate()
+
+with open("./labels.json", "w") as f:
+    json.dump(id_to_label, f)
+
+trainer.save_model("./model")