Rezuwan commited on
Commit
a86288a
·
verified ·
1 Parent(s): c79f2e7

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +29 -29
app.py CHANGED
@@ -1,29 +1,29 @@
1
- import gradio as gr
2
- import onnxruntime as rt
3
- from transformers import AutoTokenizer
4
- import torch, json
5
-
6
- tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")
7
-
8
- with open("genre_types_encoded_multi_class.json", "r") as fp:
9
- encode_genre_types = json.load(fp)
10
-
11
- genres = list(encode_genre_types.keys())
12
-
13
-
14
- inf_session = rt.InferenceSession('paper_task_classifier_quantized.onnx')
15
- input_name = inf_session.get_inputs()[0].name
16
- output_name = inf_session.get_outputs()[0].name
17
-
18
-
19
- def classify_film_genre(description):
20
- input_ids = tokenizer(description)['input_ids'][:512]
21
- logits = inf_session.run([output_name], {input_name: [input_ids]})[0]
22
- logits = torch.FloatTensor(logits)
23
- probs = torch.sigmoid(logits)[0]
24
- return dict(zip(genres, map(float, probs)))
25
-
26
- label = gr.outputs.Label(num_top_classes=5)
27
- iface = gr.Interface(fn=classify_film_genre, inputs="text", outputs=label)
28
- iface.launch(inline=False)
29
-
 
1
+ import gradio as gr
2
+ import onnxruntime as rt
3
+ from transformers import AutoTokenizer
4
+ import torch, json
5
+
6
+ tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")
7
+
8
+ with open("genre_types_encoded_multi_class.json", "r") as fp:
9
+ encode_genre_types = json.load(fp)
10
+
11
+ genres = list(encode_genre_types.keys())
12
+
13
+
14
+ inf_session = rt.InferenceSession('film_genre_classifier_quantized.onnx')
15
+ input_name = inf_session.get_inputs()[0].name
16
+ output_name = inf_session.get_outputs()[0].name
17
+
18
+
19
+ def classify_film_genre(description):
20
+ input_ids = tokenizer(description)['input_ids'][:512]
21
+ logits = inf_session.run([output_name], {input_name: [input_ids]})[0]
22
+ logits = torch.FloatTensor(logits)
23
+ probs = torch.sigmoid(logits)[0]
24
+ return dict(zip(genres, map(float, probs)))
25
+
26
+ label = gr.outputs.Label(num_top_classes=5)
27
+ iface = gr.Interface(fn=classify_film_genre, inputs="text", outputs=label)
28
+ iface.launch(inline=False)
29
+