Upload 10 files

app.py

@@ -0,0 +1,412 @@
+from typing import Tuple, Any, List
+
+import numpy as np
+import streamlit as st
+import os
+import torch
+import random
+from pathlib import Path
+from os import PathLike
+
+from peft import AutoPeftModelForSequenceClassification
+from sqlalchemy.sql.functions import random
+from transformers import AutoTokenizer, AutoConfig, TokenizersBackend, SentencePieceBackend
+from typing import Dict
+import requests
+from bs4 import BeautifulSoup
+import json
+import pandas as pd
+
+from deep_translator import (GoogleTranslator,
+                             MyMemoryTranslator,
+                             LingueeTranslator,
+                             PonsTranslator,
+                             single_detection)
+
+os.environ["TRANSFORMERS_VERBOSITY"] = "error"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+NUM_LABELS = 149
+
+
+@st.cache_resource(show_spinner="Загрузка модели...")
+def load_model(
+        model_name: str = "oracat/bert-paper-classifier-arxiv",
+        model_path: str | PathLike[str]= Path('./checkpoints/bert-paper-classifier-arxiv/checkpoints/final_model')
+) -> Tuple[Any, Any, AutoConfig]:
+
+    model = AutoPeftModelForSequenceClassification.from_pretrained(
+        model_path,
+        num_labels=149,
+        problem_type="multi_label_classification",
+        ignore_mismatched_sizes=True
+    )
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model_cfg = AutoConfig.from_pretrained(model_name)
+
+    model = model.to(device)
+    model.eval()
+    return model, tokenizer, model_cfg
+
+
+@st.cache_data(show_spinner='Получение категорий с Arxiv...')
+def fetch_arxiv_cat_describes(url: str = "https://arxiv.org/category_taxonomy") -> Dict[str, str]:
+    resp = requests.get(url)
+    resp.raise_for_status()
+    soup = BeautifulSoup(resp.text, "html5lib")
+
+    cat2descr = {}
+    for head in soup.select("h2.accordion-head"):
+
+        body = head.find_next_sibling("div", class_="accordion-body")
+        if body is None:
+            continue
+
+        for block in body.select("div.columns.divided"):
+            h4 = block.find("h4")
+            if h4 is None:
+                continue
+
+            cat_code = h4.contents[0].text.strip()
+            cat2descr[cat_code] = h4.contents[1].text.strip('()')
+
+    return cat2descr
+
+
+@st.cache_data
+def get_ids_to_cats(json_path=Path('./data/ids2cat.json')) -> Dict[int, str]:
+    if not json_path.exists():
+        raise FileNotFoundError(json_path)
+
+    with open(json_path, 'r') as f:
+        ids2cat = json.load(f)
+        ids2cat = {int(k): v for k, v in ids2cat.items()}
+
+    return ids2cat
+
+
+@st.cache_data(show_spinner=False)
+def translate_text(text: str, translator_name: str, source: str = "auto") -> str:
+    if not text or not text.strip():
+        return text
+
+    if translator_name == "Google Translator":
+        return GoogleTranslator(source=source, target='en').translate(text)
+
+    elif translator_name == "MyMemory Translator":
+        return MyMemoryTranslator(source=source, target='en-GB').translate(text)
+
+    raise ValueError(f"Неизвестный переводчик: {translator_name}")
+
+
+
+@torch.inference_mode()
+def predict(text: str) -> List[Tuple[int, float]]:
+
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        padding=True,
+        max_length=model_cfg.max_position_embeddings,
+    )
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    outputs = model(**inputs)
+    probs = torch.sigmoid(outputs.logits[0]).cpu().tolist()
+    return [(i, p) for i, p in enumerate(probs)]
+
+
+def go_to_phase1():
+    st.session_state.step = 1
+
+
+def go_to_phase2():
+    if ((st.session_state.title_input is None or len(st.session_state.title_input) == 0 )
+            and (st.session_state.abstract_input is None or len(st.session_state.abstract_input) == 0)):
+        st.toast('Пустой заголовок и текст', icon="🚫")
+        return
+
+    st.session_state.step = 2
+    # Сохранение пользовательского ввода
+    st.session_state.saved_title = st.session_state.title_input
+    st.session_state.saved_abstract = st.session_state.abstract_input
+    st.session_state.saved_k = st.session_state.k_input
+    st.session_state.saved_translator = st.session_state.translator_choice
+
+
+def clear_form():
+    st.session_state.title_input = None
+    st.session_state.abstract_input = None
+    st.session_state.k_input = 10
+    st.session_state.translator_choice = "Google Translator"
+
+    st.session_state.saved_title = None
+    st.session_state.saved_abstract = None
+    st.session_state.saved_k = 10
+    st.session_state.saved_translator = "Google Translator"
+
+
+if 'saved_title' not in st.session_state:
+    st.session_state.saved_title = None
+
+if 'saved_abstract' not in st.session_state:
+    st.session_state.saved_abstract = None
+
+if 'saved_k' not in st.session_state:
+    st.session_state.saved_k = 10
+
+if 'saved_translator' not in st.session_state:
+    st.session_state.saved_translator = "Google Translator"
+
+if 'step' not in st.session_state:
+    st.session_state.step = 1
+
+
+model, tokenizer, model_cfg = load_model()
+cat2descr = fetch_arxiv_cat_describes()
+ids2cat = get_ids_to_cats()
+
+st.markdown("""
+<style>
+html, body, [class*="css"] {
+    font-family: "Nunito Sans", sans-serif;
+}
+
+h1, h2, h3, h4, h5, h6,
+[data-testid="stMarkdownContainer"] h1,
+[data-testid="stMarkdownContainer"] h2,
+[data-testid="stMarkdownContainer"] h3 {
+    font-family: "Instrument Serif", serif !important;
+    font-weight: 400 !important;
+}
+
+code, pre, kbd, samp {
+    font-family: "JetBrains Mono", monospace !important;
+}
+</style>
+""", unsafe_allow_html=True)
+
+
+with st.container(border=True):
+    st.caption("О ПРОЕКТЕ")
+    st.subheader("BERT Multi-label классификатор arXiv-статей")
+
+    st.markdown(
+        """
+        Проект представляет собой систему **multi-label классификации научных статей arXiv**
+        по **заголовку** и **аннотации**.
+
+        Датасет был собран напрямую через **arXiv API**: статьи выгружались по категориям,
+        затем сопоставлялись с **официальной таксономией arXiv** и дедуплицировались
+        по нормализованному идентификатору. В результате был получен корпус
+        объёмом **более 1M+ примеров**.
+        """
+    )
+
+    st.markdown(
+        """
+        Модель предсказывает **сразу несколько тегов** для одной статьи.
+        Чтобы система была устойчивой к неполным данным, в обучении используются
+        два режима входа:
+
+        - **Только заголовок**
+        - **Заголовок + аннотация**
+
+        Выбор режима происходит прямо в **batch collator**, поэтому модель учится
+        работать и с полным описанием статьи, и с сокращённым вариантом.
+        """
+    )
+
+    st.markdown(
+        """
+        В качестве основы используется **BERT для sequence classification**,
+        дообучаемый с помощью **LoRA**.
+
+        Для борьбы с сильным дисбалансом классов применяется
+        **взвешенная binary cross-entropy**: вес каждого класса рассчитывается
+        по соотношению отрицательных и положительных примеров.
+        Итоговая модель развёрнута в **Streamlit-приложении**
+        для интерактивного предсказания arXiv-категорий.
+        """
+    )
+
+    st.divider()
+
+    c1, c2, c3 = st.columns(3)
+    c1.metric("Датасет", "1M+ статей")
+    c2.metric("Обучение", "~6 часов")
+    c3.metric("GPU", "H100")
+
+    st.caption(
+        "Trainable params: 999,317 · All params: 110,596,138 · Trainable%: 0.9036"
+    )
+
+st.markdown("### Гиперпараметры обучения")
+
+train_params = pd.DataFrame(
+    [
+        ("Базовая модель", "oracat/bert-paper-classifier-arxiv"),
+        ("Тип задачи", "Multi-label classification"),
+        ("Число классов", "149"),
+        ("Разделение датасета", "90% train / 10% test"),
+        ("Batch size (train)", "512"),
+        ("Batch size (eval)", "512"),
+        ("Число эпох", "10"),
+        ("Learning rate", "5e-4"),
+        ("LR scheduler", "cosine"),
+        ("Warmup steps", "10"),
+        ("Оптимизатор", "adamw_torch_fused"),
+        ("Weight decay", "0.001"),
+        ("Gradient accumulation steps", "1"),
+        ("Mixed precision", "bf16"),
+        ("LoRA target modules", "query, key, value"),
+        ("LoRA rank (r)", "16"),
+        ("LoRA alpha", "32"),
+        ("LoRA dropout", "0.05"),
+        ("use_rslora", "True"),
+        ("modules_to_save", "classifier"),
+        ("Функция потерь", "Weighted BCEWithLogitsLoss"),
+    ],
+    columns=["Гиперпараметр", "Значение"]
+)
+
+st.dataframe(train_params, use_container_width=True, hide_index=True)
+
+if st.session_state.step == 1:
+    st.subheader('Часть 1: Ввод статьи')
+    with st.form(key='insert form', height='content'):
+
+        title = st.text_input(
+            label='Введите заголовок статьи',
+            placeholder='Заголовок статьи',
+            key='title_input',
+            value=st.session_state.saved_title
+        )
+
+        abstract = st.text_area(
+            label='Введите текст статьи',
+            placeholder='Текст статьи',
+            key='abstract_input',
+            value = st.session_state.saved_abstract,
+            height='content'
+        )
+
+        slider_help = 'Параметр регулирует вывод K наиболее вероятных тэгов из таксономии arxiv для статьи '
+        k = st.slider(
+            label='Введите top-K',
+            min_value=1,
+            max_value=20,
+            value = st.session_state.saved_k,
+            key='k_input',
+            help=slider_help
+        )
+
+        options = ["Google Translator", "MyMemory Translator"]
+        selectbox_help = ('Модель была обучена понимать только английский, поэтому если текст на другом языке'
+                          ' или на смеси языков, то его необходимо сначала перевести.')
+        option = st.selectbox(
+            label='Выбор средства перевода',
+            options=options,
+            key='translator_choice',
+            help=selectbox_help
+        )
+
+        col1, col2, col3 = st.columns([1, 1, 5], gap="small")
+
+        with col1:
+            pressed_input = st.form_submit_button(
+                "Ввод",
+                width="stretch",
+                on_click=go_to_phase2
+            )
+
+        with col2:
+            pressed_clear = st.form_submit_button(
+                "Очистка",
+                width="stretch",
+                on_click=clear_form
+            )
+
+elif st.session_state.step == 2:
+
+    if np.random.rand(1) < 0.5:
+        st.balloons()
+    else:
+        st.snow()
+
+    st.subheader('Часть 2: Результаты классификации')
+
+    title = (
+        st.session_state.saved_title
+        if st.session_state.saved_title is not None
+        else ''
+    )
+
+    abstract = (
+        st.session_state.saved_abstract
+        if st.session_state.saved_abstract is not None
+        else ''
+    )
+
+    k = (
+        st.session_state.saved_k
+        if st.session_state.saved_k is not None
+        else 10
+    )
+
+    translator_name = (
+        st.session_state.saved_translator
+        if st.session_state.saved_translator is not None
+        else "Google Translator"
+    )
+
+    if title and abstract:
+        input_text = title + '\n\n' + abstract
+    elif title:
+        input_text = title
+    else:
+        input_text = abstract
+
+    lang = single_detection(input_text, api_key='db46a38fbac48c2b159384593d450933')
+    translated_text = input_text
+    if lang != 'en':
+        with st.spinner(f"Перевод через {translator_name}..."):
+            try:
+                translated_text = translate_text(
+                    text=input_text,
+                    translator_name=translator_name,
+                    source='auto'
+                )
+            except Exception as e:
+                st.error(f"Ошибка перевода: {e}")
+                st.stop()
+
+    with st.spinner("Модель анализирует текст..."):
+        probabilities = sorted(predict(translated_text), key=lambda x: x[1], reverse=True)
+        probs_k = [(ids2cat[x[0]], x[1]) for x in probabilities[:k]]
+
+    with st.container(border=True):
+        st.caption("Лучшие совпадения")
+        c1, c2 = st.columns([3, 1])
+
+        with c1:
+            best_cat = cat2descr[probs_k[0][0]]
+            st.markdown(f"# **{best_cat}**")
+            st.caption(f"Arxiv Tag: {probs_k[0][0]}")
+
+        with c2:
+            best_p = probs_k[0][1]
+            st.markdown(f"# _**{round(best_p * 100, 2)}%**_")
+
+        with st.container(border=True):
+            st.caption(f"Наиболее подходящие {k}/{NUM_LABELS} категорий")
+            st.markdown(
+                "<hr style='margin: 8px 0; border: none; border-top: 1px solid #ddd;'>",
+                unsafe_allow_html=True
+            )
+            rows = []
+            for cat, prob in probs_k:
+                rows.append({'Category': cat2descr[cat], 'Probability': f'{round(prob * 100, 2)}%', 'Arxiv Tag': cat})
+            df = pd.DataFrame(rows)
+            st.dataframe(df)
+
+        st.button(label='Назад', on_click=go_to_phase1, width="stretch")

checkpoints/bert-paper-classifier-arxiv/checkpoints/final_model/README.md

@@ -0,0 +1,206 @@
+---
+base_model: oracat/bert-paper-classifier-arxiv
+library_name: peft
+tags:
+- base_model:adapter:oracat/bert-paper-classifier-arxiv
+- lora
+- transformers
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+### Framework versions
+
+- PEFT 0.18.1

checkpoints/bert-paper-classifier-arxiv/checkpoints/final_model/adapter_config.json

@@ -0,0 +1,46 @@
+{
+  "alora_invocation_tokens": null,
+  "alpha_pattern": {},
+  "arrow_config": null,
+  "auto_mapping": null,
+  "base_model_name_or_path": "oracat/bert-paper-classifier-arxiv",
+  "bias": "none",
+  "corda_config": null,
+  "ensure_weight_tying": false,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 32,
+  "lora_bias": false,
+  "lora_dropout": 0.05,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": [
+    "classifier",
+    "classifier",
+    "score"
+  ],
+  "peft_type": "LORA",
+  "peft_version": "0.18.1",
+  "qalora_group_size": 16,
+  "r": 16,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "query",
+    "value",
+    "key"
+  ],
+  "target_parameters": null,
+  "task_type": "SEQ_CLS",
+  "trainable_token_indices": null,
+  "use_dora": false,
+  "use_qalora": false,
+  "use_rslora": true
+}

checkpoints/bert-paper-classifier-arxiv/checkpoints/final_model/adapter_model.safetensors

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

data/cat2ids.json

@@ -0,0 +1,151 @@
+{
+    "astro-ph.CO": 0,
+    "astro-ph.EP": 1,
+    "astro-ph.GA": 2,
+    "astro-ph.HE": 3,
+    "astro-ph.IM": 4,
+    "astro-ph.SR": 5,
+    "cond-mat.dis-nn": 6,
+    "cond-mat.mes-hall": 7,
+    "cond-mat.mtrl-sci": 8,
+    "cond-mat.other": 9,
+    "cond-mat.quant-gas": 10,
+    "cond-mat.soft": 11,
+    "cond-mat.stat-mech": 12,
+    "cond-mat.str-el": 13,
+    "cond-mat.supr-con": 14,
+    "cs.AI": 15,
+    "cs.AR": 16,
+    "cs.CC": 17,
+    "cs.CE": 18,
+    "cs.CG": 19,
+    "cs.CL": 20,
+    "cs.CR": 21,
+    "cs.CV": 22,
+    "cs.CY": 23,
+    "cs.DB": 24,
+    "cs.DC": 25,
+    "cs.DL": 26,
+    "cs.DM": 27,
+    "cs.DS": 28,
+    "cs.ET": 29,
+    "cs.FL": 30,
+    "cs.GL": 31,
+    "cs.GR": 32,
+    "cs.GT": 33,
+    "cs.HC": 34,
+    "cs.IR": 35,
+    "cs.IT": 36,
+    "cs.LG": 37,
+    "cs.LO": 38,
+    "cs.MA": 39,
+    "cs.MM": 40,
+    "cs.MS": 41,
+    "cs.NE": 42,
+    "cs.NI": 43,
+    "cs.OH": 44,
+    "cs.OS": 45,
+    "cs.PF": 46,
+    "cs.PL": 47,
+    "cs.RO": 48,
+    "cs.SC": 49,
+    "cs.SD": 50,
+    "cs.SE": 51,
+    "cs.SI": 52,
+    "econ.EM": 53,
+    "econ.GN": 54,
+    "econ.TH": 55,
+    "eess.AS": 56,
+    "eess.IV": 57,
+    "eess.SP": 58,
+    "eess.SY": 59,
+    "gr-qc": 60,
+    "hep-ex": 61,
+    "hep-lat": 62,
+    "hep-ph": 63,
+    "hep-th": 64,
+    "math-ph": 65,
+    "math.AC": 66,
+    "math.AG": 67,
+    "math.AP": 68,
+    "math.AT": 69,
+    "math.CA": 70,
+    "math.CO": 71,
+    "math.CT": 72,
+    "math.CV": 73,
+    "math.DG": 74,
+    "math.DS": 75,
+    "math.FA": 76,
+    "math.GM": 77,
+    "math.GN": 78,
+    "math.GR": 79,
+    "math.GT": 80,
+    "math.HO": 81,
+    "math.KT": 82,
+    "math.LO": 83,
+    "math.MG": 84,
+    "math.NA": 85,
+    "math.NT": 86,
+    "math.OA": 87,
+    "math.OC": 88,
+    "math.PR": 89,
+    "math.QA": 90,
+    "math.RA": 91,
+    "math.RT": 92,
+    "math.SG": 93,
+    "math.SP": 94,
+    "math.ST": 95,
+    "nlin.AO": 96,
+    "nlin.CD": 97,
+    "nlin.CG": 98,
+    "nlin.PS": 99,
+    "nlin.SI": 100,
+    "nucl-ex": 101,
+    "nucl-th": 102,
+    "physics.acc-ph": 103,
+    "physics.ao-ph": 104,
+    "physics.app-ph": 105,
+    "physics.atm-clus": 106,
+    "physics.atom-ph": 107,
+    "physics.bio-ph": 108,
+    "physics.chem-ph": 109,
+    "physics.class-ph": 110,
+    "physics.comp-ph": 111,
+    "physics.data-an": 112,
+    "physics.ed-ph": 113,
+    "physics.flu-dyn": 114,
+    "physics.gen-ph": 115,
+    "physics.geo-ph": 116,
+    "physics.hist-ph": 117,
+    "physics.ins-det": 118,
+    "physics.med-ph": 119,
+    "physics.optics": 120,
+    "physics.plasm-ph": 121,
+    "physics.pop-ph": 122,
+    "physics.soc-ph": 123,
+    "physics.space-ph": 124,
+    "q-bio.BM": 125,
+    "q-bio.CB": 126,
+    "q-bio.GN": 127,
+    "q-bio.MN": 128,
+    "q-bio.NC": 129,
+    "q-bio.OT": 130,
+    "q-bio.PE": 131,
+    "q-bio.QM": 132,
+    "q-bio.SC": 133,
+    "q-bio.TO": 134,
+    "q-fin.CP": 135,
+    "q-fin.GN": 136,
+    "q-fin.MF": 137,
+    "q-fin.PM": 138,
+    "q-fin.PR": 139,
+    "q-fin.RM": 140,
+    "q-fin.ST": 141,
+    "q-fin.TR": 142,
+    "quant-ph": 143,
+    "stat.AP": 144,
+    "stat.CO": 145,
+    "stat.ME": 146,
+    "stat.ML": 147,
+    "stat.OT": 148
+}

data/ids2cat.json

@@ -0,0 +1,151 @@
+{
+    "0": "astro-ph.CO",
+    "1": "astro-ph.EP",
+    "2": "astro-ph.GA",
+    "3": "astro-ph.HE",
+    "4": "astro-ph.IM",
+    "5": "astro-ph.SR",
+    "6": "cond-mat.dis-nn",
+    "7": "cond-mat.mes-hall",
+    "8": "cond-mat.mtrl-sci",
+    "9": "cond-mat.other",
+    "10": "cond-mat.quant-gas",
+    "11": "cond-mat.soft",
+    "12": "cond-mat.stat-mech",
+    "13": "cond-mat.str-el",
+    "14": "cond-mat.supr-con",
+    "15": "cs.AI",
+    "16": "cs.AR",
+    "17": "cs.CC",
+    "18": "cs.CE",
+    "19": "cs.CG",
+    "20": "cs.CL",
+    "21": "cs.CR",
+    "22": "cs.CV",
+    "23": "cs.CY",
+    "24": "cs.DB",
+    "25": "cs.DC",
+    "26": "cs.DL",
+    "27": "cs.DM",
+    "28": "cs.DS",
+    "29": "cs.ET",
+    "30": "cs.FL",
+    "31": "cs.GL",
+    "32": "cs.GR",
+    "33": "cs.GT",
+    "34": "cs.HC",
+    "35": "cs.IR",
+    "36": "cs.IT",
+    "37": "cs.LG",
+    "38": "cs.LO",
+    "39": "cs.MA",
+    "40": "cs.MM",
+    "41": "cs.MS",
+    "42": "cs.NE",
+    "43": "cs.NI",
+    "44": "cs.OH",
+    "45": "cs.OS",
+    "46": "cs.PF",
+    "47": "cs.PL",
+    "48": "cs.RO",
+    "49": "cs.SC",
+    "50": "cs.SD",
+    "51": "cs.SE",
+    "52": "cs.SI",
+    "53": "econ.EM",
+    "54": "econ.GN",
+    "55": "econ.TH",
+    "56": "eess.AS",
+    "57": "eess.IV",
+    "58": "eess.SP",
+    "59": "eess.SY",
+    "60": "gr-qc",
+    "61": "hep-ex",
+    "62": "hep-lat",
+    "63": "hep-ph",
+    "64": "hep-th",
+    "65": "math-ph",
+    "66": "math.AC",
+    "67": "math.AG",
+    "68": "math.AP",
+    "69": "math.AT",
+    "70": "math.CA",
+    "71": "math.CO",
+    "72": "math.CT",
+    "73": "math.CV",
+    "74": "math.DG",
+    "75": "math.DS",
+    "76": "math.FA",
+    "77": "math.GM",
+    "78": "math.GN",
+    "79": "math.GR",
+    "80": "math.GT",
+    "81": "math.HO",
+    "82": "math.KT",
+    "83": "math.LO",
+    "84": "math.MG",
+    "85": "math.NA",
+    "86": "math.NT",
+    "87": "math.OA",
+    "88": "math.OC",
+    "89": "math.PR",
+    "90": "math.QA",
+    "91": "math.RA",
+    "92": "math.RT",
+    "93": "math.SG",
+    "94": "math.SP",
+    "95": "math.ST",
+    "96": "nlin.AO",
+    "97": "nlin.CD",
+    "98": "nlin.CG",
+    "99": "nlin.PS",
+    "100": "nlin.SI",
+    "101": "nucl-ex",
+    "102": "nucl-th",
+    "103": "physics.acc-ph",
+    "104": "physics.ao-ph",
+    "105": "physics.app-ph",
+    "106": "physics.atm-clus",
+    "107": "physics.atom-ph",
+    "108": "physics.bio-ph",
+    "109": "physics.chem-ph",
+    "110": "physics.class-ph",
+    "111": "physics.comp-ph",
+    "112": "physics.data-an",
+    "113": "physics.ed-ph",
+    "114": "physics.flu-dyn",
+    "115": "physics.gen-ph",
+    "116": "physics.geo-ph",
+    "117": "physics.hist-ph",
+    "118": "physics.ins-det",
+    "119": "physics.med-ph",
+    "120": "physics.optics",
+    "121": "physics.plasm-ph",
+    "122": "physics.pop-ph",
+    "123": "physics.soc-ph",
+    "124": "physics.space-ph",
+    "125": "q-bio.BM",
+    "126": "q-bio.CB",
+    "127": "q-bio.GN",
+    "128": "q-bio.MN",
+    "129": "q-bio.NC",
+    "130": "q-bio.OT",
+    "131": "q-bio.PE",
+    "132": "q-bio.QM",
+    "133": "q-bio.SC",
+    "134": "q-bio.TO",
+    "135": "q-fin.CP",
+    "136": "q-fin.GN",
+    "137": "q-fin.MF",
+    "138": "q-fin.PM",
+    "139": "q-fin.PR",
+    "140": "q-fin.RM",
+    "141": "q-fin.ST",
+    "142": "q-fin.TR",
+    "143": "quant-ph",
+    "144": "stat.AP",
+    "145": "stat.CO",
+    "146": "stat.ME",
+    "147": "stat.ML",
+    "148": "stat.OT"
+}

download_data.py

@@ -0,0 +1,158 @@
+from typing import Dict
+from collections import Counter
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+import json
+import re
+from pathlib import Path
+import arxiv
+import pandas as pd
+import requests
+from bs4 import BeautifulSoup
+from tqdm.auto import tqdm
+
+
+TAXONOMY_URL = "https://arxiv.org/category_taxonomy"
+
+
+def fetch_arxiv_categories(url: str = "https://arxiv.org/category_taxonomy") -> Dict[str, str]:
+    resp = requests.get(url)
+    resp.raise_for_status()
+    soup = BeautifulSoup(resp.text, "html5lib")
+
+    cat2main = {}
+    for head in soup.select("h2.accordion-head"):
+        main_class = head.get_text(strip=True)
+
+        body = head.find_next_sibling("div", class_="accordion-body")
+        if body is None:
+            continue
+
+        for block in body.select("div.columns.divided"):
+            h4 = block.find("h4")
+            if h4 is None:
+                continue
+
+            cat_code = h4.contents[0].text.strip()
+            cat2main[cat_code] = main_class
+
+    return cat2main
+
+def normalize_arxiv_id(arxiv_id: str) -> str:
+    raw_id = arxiv_id.split("/")[-1]   # 2107.05580v1
+    raw_id = re.sub(r"v\d+$", "", raw_id)   # 2107.05580
+    raw_id = re.sub(r"\D", "", raw_id)      # 210705580
+    return raw_id
+
+def make_dataset(data_json_path, remain_cats_path, categories: Dict[str, str], num_class_samples: int) -> None:
+
+    if Path(data_json_path).exists():
+        with open(data_json_path, "r", encoding="utf-8") as f:
+            records = json.load(f)
+    else:
+        records = {}
+
+    if Path(remain_cats_path).exists():
+        with open(remain_cats_path, "r", encoding="utf-8") as f:
+            remain_cats = f.readline().split()
+    else:
+        remain_cats = list(categories.keys())
+
+    client = arxiv.Client(
+        page_size=1000,
+        delay_seconds=3,
+        num_retries=3
+    )
+
+    processed_cats = []
+    try:
+        for category in tqdm(remain_cats, 'Fetching Categories'):
+
+            search = arxiv.Search(
+                query=f"cat:{category}",
+                max_results=num_class_samples,
+                sort_by=arxiv.SortCriterion.Relevance,
+                sort_order=arxiv.SortOrder.Descending
+            )
+
+            buffer = {}
+            for result in client.results(search):
+                if any([result.title is None,
+                        result.categories is None,
+                        result.entry_id is None,
+                        result.summary is None]):
+                    continue
+
+                arxiv_id = normalize_arxiv_id(result.entry_id)
+
+                base_cats = set()
+                target_cats = set()
+                for cat in result.categories:
+                    if cat not in categories.keys():
+                        continue
+                    base_cats.add(categories[cat])
+                    target_cats.add(cat)
+
+                if arxiv_id not in records.keys():
+                    new_record = {
+                        'title': result.title,
+                        'abstract': result.summary,
+                        'categories': list(target_cats),
+                        'base_categories': list(base_cats),
+                    }
+
+                    buffer[arxiv_id] = new_record
+
+            records.update(buffer)
+            del buffer
+            processed_cats.append(category)
+
+    finally:
+        with open(data_json_path, "w", encoding="utf-8") as f:
+            json.dump(records, f, ensure_ascii=False, indent=4)
+
+        processed_cats_set = set(processed_cats)
+        new_remain_cats = [cat for cat in remain_cats if cat not in processed_cats_set]
+
+        with open(remain_cats_path, "w", encoding="utf-8") as f:
+            f.write(" ".join(new_remain_cats))
+
+
+if __name__ == "__main__":
+    data_path = Path('dataset.json')
+    remain_cats_path = Path('./logs/remain_cats.txt')
+
+    remain_cats_path.parent.mkdir(parents=True, exist_ok=True)
+    categories = fetch_arxiv_categories("https://arxiv.org/category_taxonomy")
+
+    if not remain_cats_path.exists():
+        with open(remain_cats_path, 'w', encoding='utf-8') as file:
+            file.write(' '.join(categories.keys()))
+
+    #make_dataset(data_path, remain_cats_path, categories, 2)
+
+    with open(data_path, "r", encoding="utf-8") as f:
+        dataset = json.load(f)
+
+    dataset = pd.DataFrame.from_dict(dataset, orient='index').reset_index()
+
+    data = []
+    for x in dataset['categories']:
+        data.extend(x)
+
+    cnt = Counter(data)
+    print(len(cnt.keys()))
+
+    items = cnt.most_common()
+
+    labels = [x[0] for x in items]
+    values = [x[1] for x in items]
+
+    plt.figure(figsize=(10, max(4, len(labels) * 0.4)))
+    plt.barh(labels, values)
+    plt.xlabel("Count")
+    plt.ylabel("Object")
+    plt.title("Counter result")
+    plt.gca().invert_yaxis()
+    plt.tight_layout()
+    plt.show()

inference.py

@@ -0,0 +1,74 @@
+from peft import AutoPeftModelForSequenceClassification
+from transformers import AutoTokenizer
+from pathlib import Path
+import torch
+import json
+
+model_name = "oracat/bert-paper-classifier-arxiv"
+output_model_path = Path(f"./checkpoints/{model_name.split('/')[-1]}/checkpoints")
+save_path = Path("./data")
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Загружаем id -> category
+with open(save_path / "ids2cat.json", "r", encoding="utf-8") as f:
+    ids2cat = json.load(f)
+    ids2cat = {int(k): v for k, v in ids2cat.items()}
+
+num_labels = len(ids2cat)
+
+# Загружаем PEFT-модель из папки final_model
+model = AutoPeftModelForSequenceClassification.from_pretrained(
+    output_model_path / "final_model",
+    num_labels=num_labels,
+    problem_type="multi_label_classification",
+    ignore_mismatched_sizes=True
+)
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+model = model.to(device)
+model.eval()
+
+
+@torch.inference_mode()
+def predict(
+    text: str,
+    threshold: float = 0.5,
+    max_length: int = 512,
+):
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        padding=True,
+        max_length=max_length,
+    )
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+
+    outputs = model(**inputs)
+
+    probs = torch.sigmoid(outputs.logits[0]).cpu()
+
+    pred_ids = (probs >= threshold).nonzero(as_tuple=True)[0].tolist()
+    pred_labels = [ids2cat[i] for i in pred_ids]
+
+    return {
+        "labels": pred_labels,
+        "scores": {ids2cat[i]: float(probs[i]) for i in range(num_labels)},
+    }
+
+
+if __name__ == "__main__":
+    text = """Macroscopic transport patterns of UAV traffic in 3D anisotropic wind fields: A constraint-preserving hybrid PINN-FVM approach""" + '\n\n' + """Macroscopic unmanned aerial vehicle (UAV) traffic organization in three-dimensional airspace faces significant challenges from static wind fields and complex obstacles. A critical difficulty lies in simultaneously capturing the strong anisotropy induced by wind while strictly preserving transport consistency and boundary semantics, which are often compromised in standard physics-informed learning approaches. To resolve this, we propose a constraint-preserving hybrid solver that integrates a physics-informed neural network for the anisotropic Eikonal value problem with a conservative finite-volume method for steady density transport. These components are coupled through an outer Picard iteration with under-relaxation, where the target condition is hard-encoded and strictly conservative no-flux boundaries are enforced during the transport step. We evaluate the framework on reproducible homing and point-to-point scenarios, effectively capturing value slices, induced-motion patterns, and steady density structures such as bands and bottlenecks. Ultimately, our perspective emphasizes the value of a reproducible computational framework supported by transparent empirical diagnostics to enable the traceable assessment of macroscopic traffic phenomena."""
+
+    result = predict(text, threshold=0.7)
+
+    print("Predicted labels:")
+    for label in result["labels"]:
+        print(label)
+
+    print("\nTop-10 scores:")
+    top_scores = sorted(result["scores"].items(), key=lambda x: x[1], reverse=True)[:10]
+    for label, score in top_scores:
+        print(f"{label}: {score:.4f}")

preprocess.py

@@ -0,0 +1,107 @@
+from os import PathLike
+from transformers import AutoTokenizer, PreTrainedTokenizerBase, AutoConfig
+from datasets import Dataset
+from typing import Dict, Any, Tuple, List
+from pathlib import Path
+import json
+
+
+def make_full_text_(title: str, abstract: str) -> str:
+    if abstract is None:
+        abstract = ""
+    return title + "\n\n" + abstract
+
+
+def transform_(
+    examples: Dict[str, Any],
+    tokenizer: PreTrainedTokenizerBase,
+    categories_column: str,
+    cat2ids: Dict[str, int],
+    tokenizer_cfg: Dict[str, Any],
+):
+    # batched=True => examples это dict списков
+
+    full_texts = [
+        make_full_text_(title, abstract)
+        for title, abstract in zip(examples["title"], examples["abstract"])
+    ]
+
+    title_tokens = tokenizer(examples["title"], **tokenizer_cfg)
+    full_tokens = tokenizer(full_texts, **tokenizer_cfg)
+
+    return {
+        "title_input_ids": title_tokens["input_ids"],
+        "title_attention_mask": title_tokens["attention_mask"],
+
+        "full_input_ids": full_tokens["input_ids"],
+        "full_attention_mask": full_tokens["attention_mask"],
+
+        "labels_ids": [
+            [cat2ids[cat] for cat in categories]
+            for categories in examples[categories_column]
+        ],
+    }
+
+
+def json_to_dataset_(data: Dict[str, Any], categories_column: str) -> Tuple[Dataset, List[str]]:
+    rows = []
+    cats_names = set()
+    for arxiv_id, fields in data.items():
+        row = {
+            "id": arxiv_id,
+            **fields
+        }
+        rows.append(row)
+        cats_names.update(fields[categories_column])
+
+    return Dataset.from_list(rows), list(cats_names)
+
+
+def dataset_preprocess(
+        dataset_path: str | PathLike[str],
+        classifier_name: str,
+        categories_column: str):
+
+    if not Path(dataset_path).exists():
+        raise FileNotFoundError
+
+    with open(dataset_path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    dataset, categories = json_to_dataset_(data, categories_column)
+    categories = sorted(categories)
+
+    cat2ids = {}
+    ids2cat = {}
+
+    for idx, cat in enumerate(categories):
+        cat2ids[cat] = idx
+        ids2cat[idx] = cat
+
+    tokenizer = AutoTokenizer.from_pretrained(classifier_name)
+    model_cfg = AutoConfig.from_pretrained(classifier_name)
+
+    tokenizer_cfg = {
+        "truncation": True,
+        "padding": "max_length",
+        "max_length": model_cfg.max_position_embeddings
+    }
+
+    dataset = dataset.map(
+        lambda examples: transform_(
+            examples,
+            tokenizer,
+            categories_column,
+            cat2ids,
+            tokenizer_cfg
+        ),
+        batched=True,
+        batch_size=10_000,
+        remove_columns=dataset.column_names,
+    )
+
+    return {
+        "dataset": dataset,
+        "cat2ids": cat2ids,
+        "ids2cat": ids2cat,
+    }

train.py

@@ -0,0 +1,252 @@
+import json
+from collections import Counter, defaultdict
+from os import PathLike
+from typing import Any, Dict, List
+
+from torch import nn
+from transformers import Trainer, AutoModelForSequenceClassification, TrainingArguments
+from peft import get_peft_model, LoraConfig, TaskType
+from datasets import load_from_disk
+from preprocess import dataset_preprocess
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+class BatchCollator:
+    def __init__(
+        self,
+        num_classes: int,
+        p_only_title: float = 0.4,
+    ):
+        self.num_classes = num_classes
+        self.p_only_title = p_only_title
+
+    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
+        batch_size = len(features)
+
+        use_only_title = np.random.rand(batch_size) < self.p_only_title
+
+        input_ids = []
+        attention_mask = []
+
+        for i, feat in enumerate(features):
+
+            required_keys = [
+                'title_input_ids',
+                'title_attention_mask',
+                'full_input_ids',
+                'full_attention_mask',
+                'labels_ids',
+            ]
+
+            if not all(k in feat for k in required_keys):
+                raise KeyError(f'{required_keys} must be in dataset elements!')
+
+            if use_only_title[i]:
+                input_ids.append(feat["title_input_ids"])
+                attention_mask.append(feat["title_attention_mask"])
+            else:
+                input_ids.append(feat["full_input_ids"])
+                attention_mask.append(feat["full_attention_mask"])
+
+        input_ids = torch.tensor(input_ids, dtype=torch.long)
+        attention_mask = torch.tensor(attention_mask, dtype=torch.long)
+
+        labels = torch.zeros((batch_size, self.num_classes), dtype=torch.float32)
+
+        for i, el in enumerate(features):
+            for label_id in el["labels_ids"]:
+                labels[i, label_id] = 1.0
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels,
+        }
+
+def get_dataset(
+        dataset_json_path: str | PathLike[str],
+        model_name: str,
+        categories_column: str
+) -> Dict[str, Any]:
+
+    save_path = Path('./data')
+
+    if (save_path / 'tokenized_dataset').exists():
+        dataset = load_from_disk(str(save_path / 'tokenized_dataset'))
+
+        with open(save_path / 'cat2ids.json', 'r') as f:
+            cat2ids = json.load(f)
+            cat2ids = {k: int(v) for k, v in cat2ids.items()}
+
+        with open(save_path / 'ids2cat.json', 'r') as f:
+            ids2cat = json.load(f)
+            ids2cat = {int(k): v for k, v in ids2cat.items()}
+
+        return {
+            "dataset": dataset,
+            "cat2ids": cat2ids,
+            "ids2cat": ids2cat,
+        }
+
+    result = dataset_preprocess(
+        dataset_json_path,
+        model_name,
+        categories_column
+    )
+    save_data(result, './data')
+
+    return result
+
+def save_data(result: Dict[str, Any], save_path: str | PathLike[str]):
+    save_path = Path(save_path)
+    save_path.mkdir(parents=True, exist_ok=True)
+
+    if not (save_path / 'tokenized_dataset').exists():
+        dataset = result["dataset"]
+        dataset.save_to_disk(save_path / 'tokenized_dataset')
+
+    if not (save_path / 'cat2ids.json').exists():
+        with open(save_path / 'cat2ids.json', "w", encoding="utf-8") as f:
+            json.dump(result['cat2ids'], f, ensure_ascii=False, indent=4)
+
+    if not (save_path / 'ids2cat.json').exists():
+        with open(save_path / 'ids2cat.json', "w", encoding="utf-8") as f:
+            json.dump(result['ids2cat'], f, ensure_ascii=False, indent=4)
+
+def compute_cats_weights(dataset, num_classes: int) -> torch.Tensor:
+    obj_per_cat = defaultdict(int)
+
+    for labels in dataset["labels_ids"]:
+        for label in labels:
+            obj_per_cat[label] += 1
+
+    num_obj = len(dataset)
+    pos_weights = []
+
+    for class_id in range(num_classes):
+        pos_count = obj_per_cat[class_id]
+
+        if pos_count == 0:
+            pos_weights.append(1.0)
+        else:
+            pos_weights.append((num_obj - pos_count) / pos_count)
+
+    return torch.tensor(pos_weights, dtype=torch.float32)
+
+
+def compute_loss_func(outputs, labels, num_items_in_batch=None):
+    logits = outputs.logits
+    loss = nn.functional.binary_cross_entropy_with_logits(
+        input=logits,
+        target=labels,
+        pos_weight=pos_weights.to(logits.device)
+    )
+    return loss
+
+
+if __name__ == '__main__':
+    dataset_path = Path('dataset.json')
+    model_name = 'oracat/bert-paper-classifier-arxiv'
+    output_model_path = Path(f'./checkpoints/{model_name.split("/")[-1]}/checkpoints')
+
+    categories_column = 'categories'
+
+    result = get_dataset(dataset_path, model_name, categories_column)
+
+    num_cats = len(result['cat2ids'])
+    print(f'num_cats: {num_cats}')
+
+    dataset = result['dataset']
+    dataset = dataset.train_test_split(test_size=0.1, seed=123)
+    train_dataset, test_dataset = dataset['train'], dataset['test']
+
+
+    #####
+    data = []
+    for x in train_dataset['labels_ids']:
+        data.extend(x)
+
+    cnt = Counter(data)
+    print(f'num classes in train: {len(cnt.keys())}')
+
+    data = []
+    for x in test_dataset['labels_ids']:
+        data.extend(x)
+
+    cnt = Counter(data)
+    print(f'num classes in test: {len(cnt.keys())}')
+    #####
+
+    model = AutoModelForSequenceClassification.from_pretrained(
+        model_name,
+        num_labels=num_cats,
+        problem_type='multi_label_classification',
+        ignore_mismatched_sizes=True
+    )
+
+    pos_weights = compute_cats_weights(train_dataset, num_cats)
+
+    assert pos_weights.shape[0] == num_cats
+
+    lora_config = LoraConfig(
+        modules_to_save=['classifier'],
+        use_rslora=True,
+        r=16,
+        lora_alpha=32,
+        lora_dropout=0.05,
+        task_type=TaskType.SEQ_CLS
+    )
+
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+
+    training_args = TrainingArguments(
+        per_device_train_batch_size=64,
+        per_device_eval_batch_size=64,
+
+        output_dir=output_model_path,
+        logging_dir=output_model_path / "runs",
+
+        num_train_epochs=20,
+        learning_rate=5e-4,
+        lr_scheduler_type='cosine',
+        warmup_steps=10,
+        optim='adamw_torch_fused',
+        weight_decay=0.001,
+
+        gradient_accumulation_steps=4,
+        bf16=True,
+
+        logging_strategy='epoch',
+        eval_strategy='epoch',
+
+        load_best_model_at_end=True,
+        save_only_model=False,
+        save_total_limit=2,
+        save_strategy='epoch',
+        disable_tqdm=False,
+        remove_unused_columns=False,
+        seed=42,
+        dataloader_num_workers=4
+    )
+
+    collator = BatchCollator(
+        num_classes=num_cats,
+        p_only_title=0.4
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=test_dataset,
+        data_collator=collator,
+        compute_loss_func=compute_loss_func,
+    )
+
+    trainer.train()
+
+    model.save_pretrained(output_model_path / 'final_model')