Vyke2000 commited on
Commit
eaebd8c
·
1 Parent(s): c60357e

feat: topic categorizer on spaces

Browse files
Files changed (2) hide show
  1. app.py +31 -0
  2. requirement.txt +0 -0
app.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+
3
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification
4
+
5
+ tokenizer = AutoTokenizer.from_pretrained("WebOrganizer/TopicClassifier-NoURL")
6
+ model = AutoModelForSequenceClassification.from_pretrained(
7
+ "WebOrganizer/TopicClassifier-NoURL",
8
+ trust_remote_code=True,
9
+ use_memory_efficient_attention=False)
10
+
11
+ def predict(text):
12
+ inputs = tokenizer([text], return_tensors="pt")
13
+ outputs = model(**inputs)
14
+ probs = outputs.logits.softmax(dim=-1)
15
+ pred_index = probs.argmax(dim=-1).item()
16
+ confidence_score = probs[0, pred_index]
17
+ id2label = model.config.id2label
18
+ pred_label = id2label[pred_index]
19
+
20
+ return {'topic': pred_label, 'confidence': confidence_score}
21
+
22
+ title = "URL content Topic Categorizer"
23
+
24
+ topic = gr.Interface(
25
+ fn=predict,
26
+ inputs='text',
27
+ outputs='label',
28
+ title=title,
29
+ )
30
+
31
+ topic.launch()
requirement.txt ADDED
File without changes