Updated app.py

app.py

@@ -15,20 +15,13 @@ inf_session = rt.InferenceSession('question-classifier-quantized.onnx')
 input_name = inf_session.get_inputs()[0].name
 output_name = inf_session.get_outputs()[0].name
 
-threshold = 0.5  
-
 def classify_question_tags(description):
-    input_ids = tokenizer(description)['input_ids'][:512]
-    logits = inf_session.run([output_name], {input_name: [input_ids]})[0]
-    probs = torch.sigmoid(torch.FloatTensor(logits))[0]
-    
-    filtered = {tag: float(prob) for tag, prob in zip(tags, probs) if prob > threshold}
-    
-    if not filtered:
-        topk = torch.topk(probs, k=5)
-        filtered = {tags[i]: float(probs[i]) for i in topk.indices}
-
-    return filtered 
+  input_ids = tokenizer(description)['input_ids'][:512]
+  logits = inf_session.run([output_name], {input_name: [input_ids]})[0]
+  logits = torch.FloatTensor(logits)
+  probs = torch.sigmoid(logits)[0]  
+  
+  return dict(zip(tags, map(float, probs))) 
 
 label = gr.Label(num_top_classes=5)
 iface = gr.Interface(fn=classify_question_tags, inputs="text", outputs=label)