import gradio as gr
import transformers
from transformers import AutoModelForSequenceClassification, AutoTokenizer

def classify_sentence(sent:str):
    toksentence = tokenizer(sent,truncation=True,return_tensors="pt")
    model.eval()
    with torch.no_grad():
        toksentence.to(device)
        output = model(**toksentence)
    
    return F.softmax(output.logits,dim=1).argmax(dim=1)

    
def classify_text(text:str):
    sentences = sent_tokenize(text)
    annotations = np.array(list(map(classify_sentence,sentences)),dtype=object)  
    result = list(zip(sentences,[mapping[val] for val in annotations]))
    return (annotations,result)
    
def classify_text_wrapper(text:str):
    preds,result = classify_text(text)
    n = len(preds)
    non_biased = np.where(preds==0)[0].shape[0]
    biased = np.where(preds==1)[0].shape[0]
    
    return (result,{'bias ratio':biased/n})

examples=[["[Newsoms's] obsession with masks has created an almost hostile environment in our neighborhoods and streets.\n“He won because the Election was Rigged,” Trump wrote, not referring to Biden by name, adding a list of complaints about vote counting"]]

model = AutoModelForSequenceClassification.from_pretrained("tkurtulus/autotrain-rottentomato-2981285985")
tokenizer = AutoTokenizer.from_pretrained("tkurtulus/autotrain-rottentomato-2981285985");
model.eval();

label = gr.outputs.Label(num_top_classes=None,label='')
text_h = gr.outputs.HighlightedText(color_map={'Unbiased':'#9ad1A1','Biased':'#db8a8a'},label='Classification')
inputs = gr.inputs.Textbox(placeholder=None, default="", label=None)

app = gr.Interface(fn=classify_text_wrapper,title='Bias classifier',theme='default',
                    inputs="textbox",layout='unaligned', outputs=[text_h,label], capture_session=True
                    ,examples=examples)

app.launch(inbrowser=True)