sorting

app.py

@@ -3,11 +3,12 @@ import torch
 from minicons import cwe
 from huggingface_hub import hf_hub_download
 import os
+import pandas as pd
 
 from model import FFNModule, FeatureNormPredictor, FFNParams, TrainingParams
 
 
-def predict (Word, Sentence, LLM, Norm, Layer):
+def predict (Sentence, Word, LLM, Norm, Layer):
     models = {'BERT': 'bert-base-uncased',
               'ALBERT': 'albert-xxlarge-v2',
               'roBERTa': 'roberta-base'}
@@ -33,12 +34,16 @@ def predict (Word, Sentence, LLM, Norm, Layer):
         labels = [line.rstrip() for line in file.readlines()]
 
     data = (Sentence, Word)
-    emb = lm.extract_representation(data, layer=8)
+    emb = lm.extract_representation(data, layer=Layer)
     pred = torch.nn.functional.relu(model(emb))
     pred = pred.squeeze(0)
     pred_list = pred.detach().numpy().tolist()
+
+    df = pd.DataFrame({'feature':labels, 'value':pred_list})
+    df = df.sort_values('values', ascending=False)
+    df = df[df['values'] > 0]
     
-    output = [labels[i]+'\t\t\t\t\t\t\t'+str(pred_list[i]) for i in range(len(labels)) if pred_list[i] > 0.0]
+    output = [df['features'][i]+'\t\t\t\t\t\t\t'+str(df['values'][i]) for i in range(df.shape(0))]
     return "All Positive Predicted Values:\n"+"\n".join(output)
 
 demo = gr.Interface(