prediction added

app.py

@@ -7,19 +7,30 @@ import numpy as np
 from model import FeatureNormPredictor
 
 
-def predict (word, sentence, lm_name, layer, norm):
-    if word not in sentence: return "invalid input: word not in sentence"
-    model_name = lm_name + str(layer) + '_to_' + norm
+def predict (Word, Sentence, LM, Layer, Norm):
+    if Word not in Sentence: return "invalid input: word not in sentence"
+    model_name = LM + str(Layer) + '_to_' + Norm
     lm = cwe.CWE('bert-base-uncased')
-    if layer not in range (lm.layers): return "invalid input: layer not in lm"
+    if Layer not in range (lm.layers): return "invalid input: layer not in lm"
+
     model = FeatureNormPredictor.load_from_checkpoint(
         checkpoint_path=model_name+'.ckpt',
         map_location=None
         )
     model.eval()
-    inputs = [word, sentence, lm_name, str(layer), norm]
-    outputs = [input+'\t'+str(np.random.randint(0,100, size=1)[0]) for input in inputs]
-    return "\n".join(outputs)
+
+    with open (model_name+'.txt', "r") as file:
+        labels = [line.rstrip() for line in file.readlines()]
+
+    data = (Word, Sentence)
+    embs = lm.extract_representation(data, layer=8)
+    avg = embs.sum(0)/len(data)
+    pred = torch.nn.functional.relu(model(avg))
+    pred = pred.squeeze(0)
+    pred_list = pred.detach().numpy().tolist()
+    
+    output = [labels[i]+'\t'+str(pred_list[i]) for i in range(len(labels)) if i > 0.0]
+    return "\n".join(output)
 
 demo = gr.Interface(
     fn=predict,
@@ -28,7 +39,7 @@ demo = gr.Interface(
         "text", 
         gr.Radio(["bert", "roberta", "electra"]),
         "number",
-        gr.Radio(["Binder", "McRae", "Buchanan"]),
+        gr.Radio(["binder", "mcrae", "buchanan"]),
     ],
     outputs=["text"],
 )