Spaces:
Runtime error
Runtime error
| import streamlit as st | |
| import re | |
| import torch | |
| from transformers import AlbertTokenizer, AlbertModel | |
| import pytorch_lightning as pl | |
| from huggingface_hub import hf_hub_download | |
| def download_torch_model(): | |
| model_path = hf_hub_download(repo_id="adrianmoses/hate-speech-detection", filename="pytorch_hs_model.net") | |
| print(model_path) | |
| return model_path | |
| def load_model(): | |
| model = AlbertModel.from_pretrained("albert-base-v2") | |
| return model | |
| def load_tokenizer(): | |
| tokenizer = AlbertTokenizer.from_pretrained("albert-base-v2") | |
| return tokenizer | |
| def clean_tweet(tweet): | |
| return re.sub(r'@\w+:?', "", tweet, flags=re.IGNORECASE) | |
| def tokenize(tweet): | |
| tweet = clean_tweet(tweet) | |
| tokenizer = load_tokenizer() | |
| return tokenizer(tweet, padding=True, truncation=True, max_length=64, return_tensors='pt') | |
| class HateSpeechClassifier(pl.LightningModule): | |
| def __init__(self, albert_model, dropout, hidden_dim, output_dim): | |
| super().__init__() | |
| self.model = albert_model | |
| self.l1 = torch.nn.Linear(hidden_dim, hidden_dim) | |
| self.dropout = torch.nn.Dropout(dropout) | |
| self.l2 = torch.nn.Linear(hidden_dim, output_dim) | |
| self.loss = torch.nn.NLLLoss() | |
| def forward(self, input_ids, attention_mask, token_type_ids): | |
| x = self.model(input_ids, | |
| attention_mask=attention_mask, | |
| token_type_ids=token_type_ids)[0] | |
| x = x[:, 0] | |
| x = self.dropout(torch.relu(self.l1(x))) | |
| return torch.log_softmax(self.l2(x), dim=1) | |
| def training_step(self, batch, batch_idx): | |
| input_ids, attention_masks, token_type_ids, y = batch | |
| y_hat = self(input_ids, attention_masks, token_type_ids) | |
| loss = self.loss(y_hat, y.view(-1)) | |
| return loss | |
| def validation_step(self, batch, batch_idx): | |
| input_ids, attention_masks, token_type_ids, y = batch | |
| y_hat = self(input_ids, attention_masks, token_type_ids) | |
| loss = self.loss(y_hat, y.view(-1)) | |
| return loss | |
| def configure_optimizers(self): | |
| return torch.optim.Adam(self.parameters(), lr=1e-5) | |
| def setup_model(): | |
| torch_model_path = download_torch_model() | |
| albert_model = load_model() | |
| model = HateSpeechClassifier(albert_model, 0.5, 768, 2) | |
| model.load_state_dict(torch.load(torch_model_path, map_location=torch.device('cpu'))) | |
| model.eval() | |
| return model | |
| model = setup_model() | |
| st.title("Hate Speech Detection") | |
| st.caption("Text will be truncated to 64 tokens") | |
| text = st.text_input("Enter text") | |
| encoded_input = tokenize(text) | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| input_ids = encoded_input['input_ids'] | |
| attention_mask = encoded_input['attention_mask'] | |
| token_type_ids = encoded_input['token_type_ids'] | |
| pred = model(input_ids, attention_mask, token_type_ids) | |
| print(pred) | |
| print(pred.data.max(1)) | |
| label = pred.data.max(1)[1] | |
| print(label) | |
| is_hate_speech = "YES" if label == 1 else "NO" | |
| st.write(f"Is this hate speech?: {is_hate_speech}") | |