Update app.py

app.py

@@ -92,7 +92,7 @@ def main():
         elif option == "BERT Revised":
             model = HF_BertBasedModelAppDocs.from_pretrained("GradApplicationDocuments/BERTGradRevised").to(device)
         elif option == "D-BERT Generated":
-            model = HF_BertBasedModelAppDocs.from_pretrained("GradApplicationDocuments/DistilBERTGradGen").to(device)
+            model = HF_DistilBertBasedModelAppDocs.from_pretrained("GradApplicationDocuments/DistilBERTGradGen").to(device)
         else:
             model = HF_DistilBertBasedModelAppDocs.from_pretrained("GradApplicationDocuments/DistilBERTGradRevised").to(device)
         return model
@@ -123,7 +123,7 @@ def main():
     
 
     # Load the selected trained model
-    if option in ("BERT-based model", "DistilBERT-based model (BERT light)"):
+    if option in ("BERT Generated", "BERT Revised", "D-BERT Generated", "D-BERT Revised"):
         tokenizer = load_tokenizer(option)
         model = load_model(option)
     else: