Update app.py

app.py

@@ -7,14 +7,12 @@ from prompt import Prompting
 prompting= Prompting(model=model_path)
 prompt= ". Because it was "+ prompting.tokenizer.mask_token +"."
 
-
 def predict(text):
   THRESHOLD = prompting.compute_tokens_prob(prompt, token_list1=["good"], token_list2= ["bad"])[0].item()
   res=prompting.compute_tokens_prob(text+prompt, token_list1=["good"], token_list2= ["bad"])
   if res[0] > THRESHOLD:
-    return {"POSITIVE":res[0].item()-THRESHOLD}
-  return {"NEGATIVE":THRESHOLD-res[1].item()}
-
+    return {"POSITIVE":(res[0].item()-THRESHOLD)/ (1-THRESHOLD)}
+  return {"NEGATIVE":(THRESHOLD-res[0].item())/THRESHOLD}
 
 
 import gradio as gr