Added model loading function

app.py

@@ -17,6 +17,15 @@ os.environ['WORLD_SIZE'] = "1"
 os.environ['MASTER_PORT'] = "12193"
 os.environ['MASTER_ADDR'] = "127.0.0.1"
 
+@spaces.GPU
+def load_models():
+    model = SuperposedLlama.build(ckpt_dir=weight_path, 
+                         tokenizer_path=f'{weight_path}/tokenizer.model', 
+                         max_seq_len=100, 
+                         max_batch_size=32,
+                         model_parallel_size=1)
+    return model
+
 # load_dotenv()
 # print(os.getenv("HF_ACCESS_TOKEN"))
 login(os.getenv("HF_ACCESS_TOKEN"))
@@ -36,11 +45,7 @@ n_token_sample = params["n_token_sample"]
 i_weights = params["i_weights"]
 i_length = params["i_length"]
 # Load main model
-model = SuperposedLlama.build(ckpt_dir=weight_path, 
-                         tokenizer_path=f'{weight_path}/tokenizer.model', 
-                         max_seq_len=100, 
-                         max_batch_size=32,
-                         model_parallel_size=1)
+model = load_models()
 tokenizer = Tokenizer(f'{weight_path}/tokenizer.model')
 # Create ngram models
 ngrams = make_models("ckpts-200k", bigram=True, trigram=True, fourgram=True, fivegram=True, sixgram=True, sevengram=False)