Spaces:
Sleeping
Sleeping
File size: 2,337 Bytes
2058f83 41d22ed 2058f83 41d22ed 2058f83 41d22ed 2058f83 41d22ed 2058f83 41d22ed 2058f83 41d22ed |
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 |
from transformers import AutoTokenizer, ElectraForSequenceClassification
import torch
import gradio as gr
import pickle
torch.autograd.set_grad_enabled(False)
sklearn_model = pickle.load(open('classic_pipeline.pickle', 'rb'))
model_name = "AbstractQbit/electra_large_imdb_htsplice"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = ElectraForSequenceClassification.from_pretrained(model_name)
model_reg_name = "AbstractQbit/electra_large_imdb_regression_htsplice"
model_reg = ElectraForSequenceClassification.from_pretrained(model_reg_name)
def tokenize_with_splicing(text):
tokens = tokenizer(text, truncation=False)
if len(tokens['input_ids']) > 512:
tokens['input_ids'] = tokens['input_ids'][:129] + \
[102] + tokens['input_ids'][-382:]
tokens['token_type_ids'] = [0]*512
tokens['attention_mask'] = [1]*512
return tokens
def make_stars_from_confidence(prob):
stars = round(1 + prob*9)
return '★'*stars + '☆'*(10-stars)
def make_stars_from_rating(rating):
stars = round(float(torch.clamp(rating, 1, 10)))
return '★'*stars + '☆'*(10-stars)
def run_models(review):
prob_sklearn = float(sklearn_model.predict_proba([review])[0][1])
label_sklearn = 'positive' if prob_sklearn > 0.5 else 'negative'
res = f"TF-IDF SVC trained with polarity classification thinks the review is {label_sklearn} ({100*prob_sklearn:.2f}% positive confidence).\n{make_stars_from_confidence(prob_sklearn):s}\n\n"
input = tokenize_with_splicing(review).convert_to_tensors('pt', True)
output = torch.nn.functional.softmax(model(**input).logits, dim=1)
prob_electra = float(output[0][1])
label_electra = 'positive' if prob_electra > 0.5 else 'negative'
res += f"ELECTRA trained with polarity classification thinks the review is {label_electra} ({100*prob_electra:.2f}% positive confidence).\n{make_stars_from_confidence(prob_electra):s}\n\n"
rating_electra_reg = model_reg(**input).logits[0,0]
res += f"ELECTRA trained with rating regression thinks the review is rated {rating_electra_reg:.2f}★.\n{make_stars_from_rating(rating_electra_reg):s}"
return res
demo = gr.Interface(
fn=run_models,
inputs="text",
outputs="text",
title="Movie review classification",
allow_flagging='never'
)
demo.launch()
|