Updated app.py

app.py

@@ -5,6 +5,7 @@ import torch, json
 
 tokenizer = AutoTokenizer.from_pretrained("roberta-base")
 
+
 with open("tag_types_encoded.json", "r") as fp:
   encode_tag_types = json.load(fp)
 
@@ -14,12 +15,20 @@ 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]
-  logits = torch.FloatTensor(logits)
-  probs = torch.sigmoid(logits)[0]
-  return dict(zip(tags, map(float, probs))) 
+    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 
 
 label = gr.Label(num_top_classes=5)
 iface = gr.Interface(fn=classify_question_tags, inputs="text", outputs=label)