more bert models

app.py

@@ -7,15 +7,19 @@ import os
 from model import FFNModule, FeatureNormPredictor, FFNParams, TrainingParams
 
 
-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')}
+def predict (Word, Sentence, llm, norm, layer):
+    models = {'BERT': 'bert-base-uncased',
+              'ALBERT': 'albert-xxlarge-v2',
+              'roBERTa': 'roberta-base'}
     if Word not in Sentence: return "invalid input: word not in sentence"
-    model_name = models[modelname][1]
-    lm = cwe.CWE(models[modelname][0])
+    model_name_hf = llm.lower()
+    norm_name_hf = norm.lower()
+    lm = cwe.CWE(models[llm])
 
-    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'])
+    full_name_hf = f"jwalanthi/semantic-feature-classifiers/{model_name_hf}_models_all/{model_name_hf}_to_{norm_name_hf}_layer{layer}"
+
+    model_path = hf_hub_download(f"{full_name_hf}.ckpt", use_auth_token=os.environ['TOKEN'])
+    label_path = hf_hub_download(f"{full_name_hf}.txt", use_auth_token=os.environ['TOKEN'])
 
     model = FeatureNormPredictor.load_from_checkpoint(
         checkpoint_path=model_path,
@@ -40,7 +44,9 @@ demo = gr.Interface(
     inputs=[
         "text", 
         "text", 
-        gr.Radio(["Bert Layer 8 to Binder", "Albert Layer 8 to Binder"])
+        gr.Radio(["BERT"]),
+        gr.Radio("Binder", "McRae", "Buchanan"),
+        gr.Slider(0,12)
     ],
     outputs=["text"],
 )