|
|
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(allow_output_mutation=True)
|
|
|
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 = "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"]
|
|
|
with torch.no_grad():
|
|
|
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")
|
|
|
|
