| import streamlit as st | |
| import torch | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| MODEL_NAME = 'bert-base-cased' | |
| MODEL_PATH = 'bert_model' | |
| ID2CLS = { | |
| 0: 'Computer Science', | |
| 1: 'Economics', | |
| 2: 'Electrical Engineering and Systems Science', | |
| 3: 'Mathematics', | |
| 4: 'Physics', | |
| 5: 'Quantitative Biology', | |
| 6: 'Quantitative Finance', | |
| 7: 'Statistics' | |
| } | |
| def classify(text, tokenizer, model): | |
| if not text: | |
| return [""] | |
| batch = tokenizer([text], truncation=True, padding=True, max_length=256, return_tensors="pt") | |
| outputs = model(batch['input_ids'], attention_mask=batch['attention_mask']) | |
| probabilities = torch.softmax(outputs.logits, dim=1).detach().cpu().numpy()[0] | |
| total = 0 | |
| for p in probabilities.argsort()[::-1]: | |
| field = f'{ID2CLS[p]}: {round(probabilities[p] * 100, 2)} %' | |
| st.markdown(field) | |
| total += probabilities[p] | |
| if total > 0.95: | |
| break | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
| model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=8) | |
| model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu'))) | |
| model.eval() | |
| st.markdown("## Article classifier") | |
| title = st.text_area("title") | |
| text = st.text_area("article") | |
| classify(title + text, tokenizer, model) | |