fohake commited on
Commit
c726c5a
·
verified ·
1 Parent(s): 0fd2524

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +31 -0
app.py CHANGED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModelForSequenceClassification, AutoTokenizer
3
+ import torch
4
+
5
+ # Load your model and tokenizer
6
+ model_name = "fohake/cert"
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
8
+ model = AutoModelForSequenceClassification.from_pretrained(model_name)
9
+
10
+ # Define the prediction function
11
+ def predict(text):
12
+ inputs = tokenizer(text, return_tensors="pt")
13
+ with torch.no_grad():
14
+ outputs = model(**inputs)
15
+ logits = outputs.logits
16
+ probabilities = torch.nn.functional.softmax(logits, dim=-1)
17
+ predicted_class = torch.argmax(probabilities, dim=-1).item()
18
+ confidence = probabilities[0][predicted_class].item()
19
+ return {"class": predicted_class, "confidence": confidence}
20
+
21
+ # Create the Gradio interface
22
+ iface = gr.Interface(
23
+ fn=predict,
24
+ inputs=gr.inputs.Textbox(lines=2, placeholder="Enter text here..."),
25
+ outputs="json",
26
+ title="Text Classification with CERT",
27
+ description="Enter a piece of text to classify it using the CERT model."
28
+ )
29
+
30
+ if __name__ == "__main__":
31
+ iface.launch()