import torch
from transformers import AutoTokenizer,AutoModelForSequenceClassification,Trainer,TrainingArguments
import torch.nn.functional as F
from datasets import load_dataset
import gradio as gr
model_name = 'thanhcong2001/Criminal'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name,num_labels = 2)
# Set up device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
def pre_sentiment(sentence):
    model.eval()
    inputs = tokenizer(sentence,return_tensors='pt',padding=True,truncation=True,max_length=64)
    inputs = {k:v.to(device) for k,v in inputs.items()}
    with torch.no_grad():
        outputs = model(**inputs)
        probs = F.softmax(outputs.logits,dim=-1)
    pre = torch.argmax(probs,dim=-1).item()
    score = torch.max(probs).item()
    result = 'CRIMINAL' if pre == 1 else 'NOT CRIMINAL'
    return f'{result} - Score: {score}'    
demo = gr.Interface(fn=pre_sentiment, inputs="text", outputs="text", title="Sentiment Analysis")
demo.launch()