Upload 2 files

app.py

@@ -0,0 +1,98 @@
+import streamlit as st
+import torch
+import os
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+# from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
+
+
+st.markdown('## Классификатор статей')
+st.write('Данный сервис предназначен для выбора темы статьи,     \n' \
+    'основываясь на ее названии и краткой выжимкой текста статьи.       \n' \
+    'Сервис работает благодаря fine-tune версии модели distil bert.     \n' \
+    'Данные для обучения были взяты [отсюда](https://www.kaggle.com/datasets/neelshah18/arxivdataset).       \n' \
+    'Поддерживается ввод только английского языка.')
+st.markdown('#### Введите название статьи и ее краткое содержание:')
+
+device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
+
+def create_model_and_optimizer(model, lr=1e-4, beta1=0.9, beta2=0.999, device=device):
+    model = model.to(device)
+    
+    params = []
+    for param in model.parameters():
+        if param.requires_grad:
+            params.append(param)
+    
+    optimizer = torch.optim.Adam(params, lr, [beta1, beta2])
+    return model, optimizer
+
+def title_summary_transform(items, tokenizer):
+    return tokenizer(
+        items['title'] + '[SEP]' + items['summary'],
+        padding="max_length", 
+        truncation=True
+    )
+
+def predict_category(case, model, tokenizer):
+    input_ = {key: torch.tensor(val).unsqueeze(0).to(device) for key, val in title_summary_transform(case, tokenizer).items()}
+
+    pred = []
+    pred_prob = []
+
+    with torch.no_grad():
+        logits = model(**input_).logits[0]
+        probs = torch.nn.functional.softmax(logits, dim=-1)
+        probs, indices = torch.sort(probs, descending=True)
+
+        sum_prob = 0
+        for i, prob_ in enumerate(probs):
+            pred.append(indices[i].item())
+            pred_prob.append(prob_.item())
+            sum_prob += prob_
+
+            if sum_prob > 0.95:
+                break
+    return pred, pred_prob
+
+@st.cache_resource  # кэширование
+def load_model():
+    chkp_folder = '.'
+    model_name = 'model'
+    cat_count = 358
+
+    checkpoint = torch.load(os.path.join(chkp_folder, f"{model_name}.pt"), weights_only=False)
+
+    # Создаём те же классы, что и внутри чекпоинта
+    device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
+
+    model_ = AutoModelForSequenceClassification.from_pretrained('distilbert/distilbert-base-cased', num_labels=cat_count).to(device)
+
+    for param in model_.distilbert.parameters():
+        param.requires_grad = False
+
+    for i in range(4, 6):
+        for param in model_.distilbert.transformer.layer[i].parameters():
+            param.requires_grad = True
+
+    model, optimizer = create_model_and_optimizer(model_)
+
+    # Загружаем состояния из чекпоинта
+    model.load_state_dict(checkpoint['model_state_dict'])
+    ind_to_cat = checkpoint['ind_to_cat']
+    tokenizer = AutoTokenizer.from_pretrained('distilbert/distilbert-base-cased')
+    return model, tokenizer, ind_to_cat
+
+model, tokenizer, ind_to_cat = load_model()
+
+
+case_ = {}
+
+case_['title'] = st.text_area("Название статьи:", value="")
+case_['summary'] = st.text_area("Краткое содержание:", value="")
+
+
+if case_['title'] or case_['summary']:
+    categories, probabilities = predict_category(case_, model, tokenizer)
+    st.write('Возможные категории:')
+    for i, cat in enumerate(categories):
+        st.write(f'{ind_to_cat[cat]}')

model.pt

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