feat(main): init train and dvc

.dvc/.gitignore

@@ -0,0 +1,3 @@
+/config.local
+/tmp
+/cache

.dvc/config

@@ -0,0 +1,9 @@
+[core]
+    remote = ya-s3
+    autostage = true
+['remote "ya-s3"']
+    url = s3://shad-ml-2-hw-5/dvc
+    endpointurl = https://storage.yandexcloud.net
+
+[cache]
+    type = reflink,hardlink,symlink,copy

.dvcignore

@@ -0,0 +1,3 @@
+# Add patterns of files dvc should ignore, which could improve
+# the performance. Learn more at
+# https://dvc.org/doc/user-guide/dvcignore

.idea/jsonSchemas.xml

@@ -0,0 +1,25 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="JsonSchemaMappingsProjectConfiguration">
+    <state>
+      <map>
+        <entry key="dvc.yaml">
+          <value>
+            <SchemaInfo>
+              <option name="name" value="dvc.yaml" />
+              <option name="relativePathToSchema" value="https://raw.githubusercontent.com/iterative/dvcyaml-schema/master/schema.json" />
+              <option name="applicationDefined" value="true" />
+              <option name="patterns">
+                <list>
+                  <Item>
+                    <option name="path" value="dvc.yaml" />
+                  </Item>
+                </list>
+              </option>
+            </SchemaInfo>
+          </value>
+        </entry>
+      </map>
+    </state>
+  </component>
+</project>

data/models/.gitignore

@@ -0,0 +1 @@
+/model.torch

data/raw/.gitignore

@@ -0,0 +1 @@
+/arxivData.json

data/raw/arxivData.json.dvc

@@ -0,0 +1,4 @@
+outs:
+- md5: a314e2f4eab544a46e6f95802ecde647
+  size: 72422946
+  path: arxivData.json

dvc.lock

@@ -0,0 +1,24 @@
+schema: '2.0'
+stages:
+  train:
+    cmd: poetry run train
+    deps:
+    - path: data/raw/arxivData.json
+      md5: a314e2f4eab544a46e6f95802ecde647
+      size: 72422946
+    - path: shad_mlops_transformers/model.py
+      md5: 9b932a6cb0cb46fc7c656e7c80c442e0
+      size: 2008
+      isexec: true
+    - path: shad_mlops_transformers/trainer.py
+      md5: 61acf28399fadfd2495dc48242c594ba
+      size: 3650
+    params:
+      shad_mlops_transformers/config.py:
+        Config.batch_size: 32
+        Config.random_seed: 42
+        Config.test_size: 0.2
+    outs:
+    - path: data/models/model.torch
+      md5: f110836b7b7585efdbfcb8ab7d5df76c
+      size: 438187413

dvc.yaml

@@ -0,0 +1,16 @@
+stages:
+  train:
+    cmd: poetry run train
+    deps:
+      - shad_mlops_transformers/trainer.py
+      - shad_mlops_transformers/model.py
+      - data/raw/arxivData.json
+    params:
+      - shad_mlops_transformers/config.py:
+        - Config.batch_size
+        - Config.random_seed
+        - Config.test_size
+    outs:
+      # NOTE должно совпадать с конфигом
+      - data/models/model.torch
+

pyproject.toml

@@ -12,6 +12,11 @@ python = "^3.10"
 streamlit = "^1.21.0"
 torch = "^1.13"
 transformers = "^4.27.4"
+pydantic = "^1.10.7"
+scikit-learn = "^1.2.2"
+numpy = "^1.24.2"
+loguru = "^0.7.0"
+dvc = {version = "^2.54.0", extras = ["s3"]}
 
 
 [tool.poetry.group.dev.dependencies]
@@ -20,6 +25,9 @@ black = "^23.3.0"
 docformatter = "^1.6.0"
 isort = "^5.12.0"
 
+[tool.poetry.scripts]
+train = "shad_mlops_transformers.trainer:main"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

shad_mlops_transformers/config.py

@@ -0,0 +1,17 @@
+from pathlib import Path
+
+from pydantic import BaseSettings
+
+basedir = Path(__file__).parent
+
+
+class Config(BaseSettings):
+    data_dir: Path = basedir.parent / "data"
+    raw_data_dir: Path = data_dir / "raw"
+    batch_size: int = 32
+    random_seed: int = 42
+    test_size: float = 0.2
+    weights_path: Path = data_dir / "models" / "model.torch"
+
+
+config = Config()

shad_mlops_transformers/model.py

@@ -1,8 +1,56 @@
-from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline
-
-tokenizer = AutoTokenizer.from_pretrained("Davlan/distilbert-base-multilingual-cased-ner-hrl")
-model = AutoModelForTokenClassification.from_pretrained("Davlan/distilbert-base-multilingual-cased-ner-hrl")
-nlp = pipeline("ner", model=model, tokenizer=tokenizer)
-example = "Nader Jokhadar had given Syria the lead with a well-struck header in the seventh minute."
-ner_results = nlp(example)
-print(ner_results)
+from collections import OrderedDict
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from transformers import AutoModel, AutoTokenizer
+
+from shad_mlops_transformers.config import config
+
+# example = ["Nader Jokhadar had given Syria the lead with a well-struck header in the seventh minute."]
+# model_name = "bert-base-uncased"
+# model_name = "Davlan/distilbert-base-multilingual-cased-ner-hrl"
+# tokenizer = AutoTokenizer.from_pretrained(model_name)
+# model = AutoModel.from_pretrained(model_name)
+# nlp = pipeline("ner", model=model, tokenizer=tokenizer)
+# toks = tokenizer(example, padding=True, truncation=True, return_tensors="pt")
+# with torch.no_grad():
+#     p = model(**toks)
+# print(p)
+
+
+class DocumentClassifier(nn.Module):
+    def __init__(self, n_classes: int = 2):
+        super().__init__()
+        self.model_name = "bert-base-uncased"
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.encoder = AutoModel.from_pretrained(self.model_name)
+        self.n_classes = n_classes
+        self.model = nn.Sequential(
+            OrderedDict(
+                [
+                    ("fc", nn.Linear(in_features=self.encoder.pooler.dense.out_features, out_features=n_classes)),
+                    ("sm", nn.Softmax()),
+                ]
+            )
+        )
+        self.trainable_params = self.model.parameters()
+
+    def forward(self, text):
+        tok_info = self.tokenize(text)
+        with torch.no_grad():
+            embeddings = self.encoder(**tok_info)["pooler_output"]
+        return self.model(embeddings)
+
+    def tokenize(self, x: str) -> dict:
+        return self.tokenizer(x, padding=True, truncation=True, return_tensors="pt")
+
+    def from_file(self, path: Path = config.weights_path) -> "DocumentClassifier":
+        self.load_state_dict(torch.load(path))
+        return self
+
+
+if __name__ == "__main__":
+    data = ["This article describes machine learning"]
+    model = DocumentClassifier(n_classes=61).from_file()
+    model(data)

shad_mlops_transformers/trainer.py

@@ -0,0 +1,103 @@
+import json
+
+import numpy as np
+import torch
+import torch.nn as nn
+from loguru import logger
+from sklearn.model_selection import train_test_split
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+from shad_mlops_transformers.config import config
+from shad_mlops_transformers.model import DocumentClassifier
+
+
+class ArxivDataset(Dataset):
+    def __init__(self, raw_data: list[dict]):
+        """Разово вычитываем и сохраняем весь датасет."""
+        logger.info("reading data")
+        self.x = []
+        self.y = []
+        # self.data = []
+        whitelist_labels = ["math", "cs"]
+        i = 0
+        self.class_mapper = {}
+        for item in raw_data:
+            tmp_y = []
+            # да простят мне это потомки, но там зачем-то люди засунули питоновский dict в строку!
+            for tag_desc in eval(item["tag"].replace("'", '"')):
+                real_tag: str = tag_desc["term"]
+                # пока берем только теги из whitelist
+                if not any([real_tag.startswith(x) for x in whitelist_labels]):
+                    continue
+                if real_tag not in self.class_mapper:
+                    self.class_mapper[real_tag] = i
+                    i += 1
+                tmp_y.append(self.class_mapper[real_tag])
+                # берем только один тег
+                break
+            # если был хотя бы один валидный тег, добавляем в датасет
+            if len(tmp_y):
+                # NOTE берем только один тег
+                # self.data.append({"label": tmp_y[0], "text": item["summary"]})
+                self.x.append(item["summary"])
+                self.y.append(tmp_y[0])
+        self.classes = sorted(list(self.class_mapper.keys()))
+        logger.info("[Done] reading data")
+
+    def __getitem__(self, i):
+        # return self.data[i]
+        return self.x[i], self.y[i]
+
+    def __len__(self):
+        # return len(self.data)
+        return len(self.x)
+
+
+def make_train_val():
+    with open(config.raw_data_dir / "arxivData.json", "r") as f:
+        _raw_json = json.load(f)
+    return train_test_split(_raw_json, test_size=config.test_size, shuffle=True, random_state=config.random_seed)
+
+
+def train_model(model: DocumentClassifier, optimizer: torch.optim.Optimizer, loader: DataLoader, criterion):
+    model.train()
+    losses_tr = []
+    for text, true_label in tqdm(loader):
+        optimizer.zero_grad()
+        pred = model(text)
+        loss = criterion(pred, true_label)
+
+        loss.backward()
+        optimizer.step()
+        losses_tr.append(loss.item())
+        break
+
+    return model, optimizer, np.mean(losses_tr)
+
+
+def collator(x):
+    return x[0]
+
+
+def save_model(model: DocumentClassifier):
+    config.weights_path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(model.state_dict(), config.weights_path)
+
+
+def main():
+    train, val = make_train_val()
+    dataset_train = ArxivDataset(train)
+    dataset_val = ArxivDataset(val)
+    loader_train = DataLoader(dataset_train, batch_size=config.batch_size, shuffle=True, drop_last=True)
+    loader_val = DataLoader(dataset_val, batch_size=config.batch_size, shuffle=True, drop_last=True)
+
+    model = DocumentClassifier(n_classes=len(dataset_train.classes))
+    optimizer = torch.optim.Adam(model.trainable_params)
+    loss = nn.CrossEntropyLoss()
+    train_model(model=model, optimizer=optimizer, loader=loader_train, criterion=loss)
+    save_model(model)
+
+
+if __name__ == "__main__":
+    main()