File size: 6,024 Bytes
c2c8638 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
#!/usr/bin/python3
import streamlit as st
import json
import numpy as np
import torch
from transformers import (
DebertaV2Config,
DebertaV2Model,
DebertaV2Tokenizer,
)
model_name = "microsoft/deberta-v3-base"
tokenizer = DebertaV2Tokenizer.from_pretrained(model_name)
def preprocess_text(text, tokenizer, max_length=512):
inputs = tokenizer(
text,
padding="max_length",
truncation=True,
max_length=max_length,
return_tensors="pt"
)
return inputs
def classify_text(text, model, tokenizer, device, threshold=0.5):
inputs = preprocess_text(text, tokenizer)
input_ids = inputs["input_ids"].to(device)
attention_mask = inputs["attention_mask"].to(device)
model.eval()
with torch.no_grad():
logits = model(input_ids, attention_mask)
probs = torch.sigmoid(logits)
predictions = (probs > threshold).int().numpy()
return probs.numpy(), predictions
def get_themes(text, model, tokenizer, label_to_theme, device, limit=5):
probabilities, _ = classify_text(text, model, tokenizer, device)
probabilities = probabilities / probabilities.sum()
themes = []
for label in probabilities[0].argsort()[-limit:]:
themes.append((label_to_theme[str(label)], probabilities[0][label]))
return themes
class DebertPaperClassifier(torch.nn.Module):
def __init__(self, num_labels, device, dropout_rate=0.1, class_weights=None):
super().__init__()
self.config = DebertaV2Config.from_pretrained(model_name)
self.deberta = DebertaV2Model.from_pretrained(model_name, config=self.config)
self.classifier = torch.nn.Sequential(
torch.nn.Dropout(dropout_rate),
torch.nn.Linear(self.config.hidden_size, 512),
torch.nn.LayerNorm(512),
torch.nn.GELU(),
torch.nn.Dropout(dropout_rate),
torch.nn.Linear(512, num_labels)
)
self._init_weights()
if class_weights is not None:
self.loss_fct = torch.nn.BCEWithLogitsLoss(weight=class_weights.to(device))
else:
self.loss_fct = torch.nn.BCEWithLogitsLoss()
class DebertPaperClassifierV5(torch.nn.Module):
def __init__(self, device, num_labels=47, dropout_rate=0.1, class_weights=None):
super().__init__()
self.config = DebertaV2Config.from_pretrained("microsoft/deberta-v3-base")
self.deberta = DebertaV2Model.from_pretrained("microsoft/deberta-v3-base", config=self.config)
self.classifier = torch.nn.Sequential(
torch.nn.Dropout(dropout_rate),
torch.nn.Linear(self.config.hidden_size, 512),
torch.nn.LayerNorm(512),
torch.nn.GELU(),
torch.nn.Dropout(dropout_rate),
torch.nn.Linear(512, num_labels)
)
if class_weights is not None:
self.loss_fct = torch.nn.BCEWithLogitsLoss(weight=class_weights.to(device))
else:
self.loss_fct = torch.nn.BCEWithLogitsLoss()
def forward(self, input_ids, attention_mask, labels=None):
outputs = self.deberta(
input_ids=input_ids,
attention_mask=attention_mask
)
logits = self.classifier(outputs.last_hidden_state[:, 0, :])
loss = None
if labels is not None:
loss = self.loss_fct(logits, labels)
return (loss, logits) if loss is not None else logits
def _init_weights(self):
for module in self.classifier.modules():
if isinstance(module, torch.nn.Linear):
module.weight.data.normal_(mean=0.0, std=0.02)
if module.bias is not None:
module.bias.data.zero_()
def forward(self,
input_ids,
attention_mask,
labels=None,
):
outputs = self.deberta(
input_ids=input_ids,
attention_mask=attention_mask
)
cls_output = outputs.last_hidden_state[:, 0, :]
logits = self.classifier(cls_output)
loss = None
if labels is not None:
loss = self.loss_fct(logits, labels)
return (loss, logits) if loss is not None else logits
@st.cache_resource
def load_model():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
with open('model_info/label_to_theme.json', 'r') as f:
label_to_theme = json.load(f)
model = DebertPaperClassifierV5(device=device, num_labels=len(label_to_theme)).to(device)
model.load_state_dict(torch.load("model_info/deberta_v3.pth", map_location=device))
return model, tokenizer, label_to_theme, device
def kek():
st.title("arXiv Paper Classifier")
st.markdown("""
<style>
.image-row {
display: flex;
flex-direction: row;
gap: 10px;
}
</style>
<div class="image-row">
<img width=100px src='https://storage.yandexcloud.net/lms-vault/media/cache/c9/a7/c9a754ba1b2bb5b34e1f178d4ec26f24.jpg'>
<img width=300px src='https://pic.rutubelist.ru/video/ba/b6/bab6ab515c15837e28eb6c99df192cae.jpg'>
</div>
""", unsafe_allow_html=True)
st.write("write the title or abstract to classify topic theme")
title = st.text_input("title")
abstract = st.text_area("abstract")
lim = int(st.number_input("top ? themes"))
if st.button("CLASSIFY"):
if not title and not abstract:
st.warning("empty abstract!!!")
return
text = f"{title}\n\n{abstract}" if title and abstract else title or abstract
model, tokenizer, label_to_theme, device = load_model()
with st.spinner("classifying..."):
themes = get_themes(text, model, tokenizer, label_to_theme, device, lim)
co = 0
st.success(f"top {int(lim)} results:")
for th, pr in themes:
st.write(f"{lim - co}. - {th}: {pr:.1%}")
co += 1
if __name__ == "__main__":
kek()
|