AliArshad commited on
Commit
c692bf5
1 Parent(s): 4928ce0

Create severitypredictor.py

Browse files
Files changed (1) hide show
  1. severitypredictor.py +39 -0
severitypredictor.py ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import XLNetTokenizer, XLNetForSequenceClassification
3
+ import gradio as gr
4
+
5
+ # Path to the saved model
6
+ model_path = '/content/drive/My Drive/XLNet_model_project_Core.pt' # Update with your model path
7
+
8
+ # Load the saved model
9
+ tokenizer = XLNetTokenizer.from_pretrained('xlnet-base-cased')
10
+ model = XLNetForSequenceClassification.from_pretrained('xlnet-base-cased', num_labels=2)
11
+ model.load_state_dict(torch.load(model_path))
12
+ model.eval()
13
+
14
+ # Function for prediction
15
+ def xl_net_predict(text):
16
+ inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=100)
17
+ with torch.no_grad():
18
+ outputs = model(**inputs)
19
+ logits = outputs.logits
20
+ probabilities = torch.softmax(logits, dim=1)
21
+ predicted_class = torch.argmax(probabilities).item()
22
+ return "Severe" if predicted_class == 1 else "Non-severe"
23
+
24
+ # Customizing the interface
25
+ iface = gr.Interface(
26
+ fn=xl_net_predict,
27
+ inputs=gr.Textbox(lines=2, label="Summary", placeholder="Enter text here..."),
28
+ outputs=gr.Textbox(label="Predicted Severity"),
29
+ title="XLNet Based Bug Report Severity Prediction",
30
+ description="Enter text and predict its severity (Severe or Non-severe).",
31
+ theme="huggingface",
32
+ examples=[
33
+ ["Can't open multiple bookmarks at once from the bookmarks sidebar using the context menu"],
34
+ ["Minor enhancements to make-source-package.sh"]
35
+ ],
36
+ allow_flagging=False
37
+ )
38
+
39
+ iface.launch()