Update app.py

app.py

@@ -26,11 +26,11 @@ def load_models(model_pair):
         
         # Load auxiliary model
         aux_tokenizer = AutoTokenizer.from_pretrained(aux_model_name)
-        aux_model = AutoModelForCausalLM.from_pretrained(aux_model_name).to(device)
+        aux_model = AutoModelForCausalLM.from_pretrained(aux_model_name, resume_download = True).to(device)
         
         # Load main model
         main_tokenizer = AutoTokenizer.from_pretrained(main_model_name)
-        main_model = AutoModelForCausalLM.from_pretrained(main_model_name).to(device)
+        main_model = AutoModelForCausalLM.from_pretrained(main_model_name, resume_download=True).to(device)
         
         # Create processor
         processor = TokenSwapProcessor(aux_model, main_tokenizer, aux_tokenizer=aux_tokenizer)