Rohith18SFDC commited on
Commit
bd4e1b1
·
verified ·
1 Parent(s): e353580

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +27 -0
app.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification
3
+ import torch
4
+
5
+ # Load the pre-trained model and tokenizer
6
+ model_name = "distilbert-base-uncased"
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
8
+ model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=5) # Adjust num_labels
9
+
10
+ # Define the function to get article suggestions
11
+ def suggest_articles(case_details):
12
+ inputs = tokenizer(case_details, return_tensors="pt")
13
+ with torch.no_grad():
14
+ outputs = model(**inputs)
15
+ prediction = outputs.logits.argmax(dim=1).item()
16
+ return f"Suggested Article ID: {prediction}"
17
+
18
+ # Build the Gradio interface
19
+ interface = gr.Interface(
20
+ fn=suggest_articles,
21
+ inputs="text",
22
+ outputs="text",
23
+ title="Knowledge Article Suggestion",
24
+ description="Enter case details to get relevant article suggestions."
25
+ )
26
+
27
+ interface.launch()