thanhcong2001 commited on
Commit
36f050b
·
verified ·
1 Parent(s): 709691e

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +24 -0
app.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import AutoTokenizer,AutoModelForSequenceClassification,Trainer,TrainingArguments
3
+ import torch.nn.functional as F
4
+ from datasets import load_dataset
5
+ import gradio as gr
6
+ model_name = 'thanhcong2001/Criminal'
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
8
+ model = AutoModelForSequenceClassification.from_pretrained(model_name,num_labels = 2)
9
+ # Set up device
10
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
11
+ model.to(device)
12
+ def pre_sentiment(sentence):
13
+ model.eval()
14
+ inputs = tokenizer(sentence,return_tensors='pt',padding=True,truncation=True,max_length=64)
15
+ inputs = {k:v.to(device) for k,v in inputs.items()}
16
+ with torch.no_grad():
17
+ outputs = model(**inputs)
18
+ probs = F.softmax(outputs.logits,dim=-1)
19
+ pre = torch.argmax(probs,dim=-1).item()
20
+ score = torch.max(probs).item()
21
+ result = 'CRIMINAL' if pre == 1 else 'NOT CRIMINAL'
22
+ return f'{result} - Score: {score}'
23
+ demo = gr.Interface(fn=pre_sentiment, inputs="text", outputs="text", title="Sentiment Analysis")
24
+ demo.launch()