| import os.path |
|
|
| import torch |
| import torch.nn as nn |
| from transformers import RobertaTokenizerFast, RobertaForMaskedLM |
| import streamlit as st |
|
|
|
|
| class SimpleClassifier(nn.Module): |
| def __init__(self, in_features: int, hidden_features: int, |
| out_features: int, activation=nn.ReLU()): |
| super().__init__() |
| self.bn = nn.BatchNorm1d(in_features) |
| self.in2hid = nn.Linear(in_features, hidden_features) |
| self.activation = activation |
| self.hid2hid = nn.Linear(hidden_features, hidden_features) |
| self.hid2out = nn.Linear(hidden_features, out_features) |
|
|
|
|
| |
| self.bn2 = nn.BatchNorm1d(hidden_features) |
|
|
| def forward(self, X): |
| X = self.bn(X) |
| X = self.in2hid(X) |
|
|
| X = self.activation(X) |
| X = self.hid2hid(torch.concat((X,), 1)) |
|
|
| X = self.activation(X) |
| X = self.hid2out(torch.concat((X,), 1)) |
|
|
| X = nn.functional.sigmoid(X) |
| return X |
|
|
|
|
| @st.cache_data() |
| def load_models(): |
| model = RobertaForMaskedLM.from_pretrained("roberta-base") |
| model.lm_head = nn.Identity() |
| tokenizer = RobertaTokenizerFast.from_pretrained("roberta-base") |
| my_classifier = SimpleClassifier(768, 768, 1) |
| weights_path = os.path.join(__file__, "..", "twitter_model_91_5-.pth") |
| my_classifier.load_state_dict(torch.load(weights_path, map_location=device)) |
| my_classifier.eval() |
| return { |
| "tokenizer": tokenizer, |
| "model": model, |
| "classifier": my_classifier |
| } |
|
|
|
|
| def classify_text(text: str) -> float: |
| models = load_models() |
| tokenizer, model, classifier = models["tokenizer"], models["model"], models["classifier"] |
|
|
| X = tokenizer( |
| text, |
| truncation=True, |
| max_length=128, |
| return_tensors='pt' |
| )["input_ids"] |
|
|
| X = model.forward(X)[-1][0].sum(axis=0)[None, :] |
| return classifier(X) |
|
|
|
|
| device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") |
|
|
