rajeshradhakrishnan commited on
Commit
e658b02
1 Parent(s): 662aba8

fixed the return with probs

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -13,9 +13,9 @@ def predict(news_headline):
13
 
14
  probabilities = learner.predict(news_headline)
15
 
16
- values, indices = torch.topk(probabilities[0]['probs'], 5)
17
 
18
- return {LABELS[i]: v.item() for i, v in zip(indices, values)}
19
 
20
  interface = gr.Interface(
21
  predict,
 
13
 
14
  probabilities = learner.predict(news_headline)
15
 
16
+ # values, indices = torch.topk(probabilities[0]['probs'], 5)
17
 
18
+ return {LABELS[i]: probabilities[0]['probs'][i] for i, v in range(len(LABELS))}
19
 
20
  interface = gr.Interface(
21
  predict,