bert and albert

app.py

@@ -7,19 +7,16 @@ import os
 from model import FFNModule, FeatureNormPredictor, FFNParams, TrainingParams
 
 
-def predict (Word, Sentence, LM, Layer, Norm):
-    models = {'bert': 'bert-base-uncased'}
+def predict (Word, Sentence, modelname):
+    models = {'Bert Layer 8 to Binder': ('bert-base-uncased', 'bert8_to_binder'),
+              'Albert Layer 8 to Binder': ('albert-xxlarge-v2', 'albert8_to_binder_opt_stop')}
     if Word not in Sentence: return "invalid input: word not in sentence"
-    model_name = LM + str(Layer) + '_to_' + Norm
-    lm = cwe.CWE(models[LM])
-    if Layer not in range (lm.layers): return "invalid input: layer not in lm"
+    model_name = models[modelname][1]
+    lm = cwe.CWE(models[modelname][0])
 
-    model_path = hf_hub_download("jwalanthi/bert_layer8_to_binder", model_name+".ckpt", use_auth_token=os.environ['TOKEN'])
-    label_path = hf_hub_download("jwalanthi/bert_layer8_to_binder", model_name+".txt", use_auth_token=os.environ['TOKEN'])
+    model_path = hf_hub_download("jwalanthi/semantic-feature-classifiers", model_name+".ckpt", use_auth_token=os.environ['TOKEN'])
+    label_path = hf_hub_download("jwalanthi/semantic-feature-classifiers", model_name+".txt", use_auth_token=os.environ['TOKEN'])
 
-    # labels = "These are some fake features".split(" ")
-    # vals = np.random.randint(-10,10,(5))
-    # return model_name+" \n"+"\n".join([labels[i]+" "+str(vals[i]) for i in range(len(labels)) if vals[i]>0])
     model = FeatureNormPredictor.load_from_checkpoint(
         checkpoint_path=model_path,
         map_location=None
@@ -43,14 +40,12 @@ demo = gr.Interface(
     inputs=[
         "text", 
         "text", 
-        gr.Radio(["bert", "albert"]),
-        "number",
-        gr.Radio(["binder", "mcrae", "buchanan"]),
+        gr.Radio(["Bert Layer 8 to Binder", "Albert Layer 8 to Binder"])
     ],
     outputs=["text"],
 )
 
-demo.launch()
+demo.launch(share=True)
 
 if __name__ == "__main__":
     demo.launch(share=True)