Update app.py

app.py

@@ -68,26 +68,32 @@ def generator(input_ids, generation_config, max_new_tokens):
     return generation_output
     
 def loadModel():
-    from llama_rope_scaled_monkey_patch import replace_llama_rope_with_scaled_rope
-    replace_llama_rope_with_scaled_rope()
-    t_model = transformers.AutoModelForCausalLM.from_pretrained(
-                base_model,
-                torch_dtype=torch.float16,
-                cache_dir=cache_dir,
-                device_map="auto",
-            )
-    print(t_model.state_dict().keys())
-    model = PeftModel.from_pretrained(
-                t_model,
-                lora_weights,
-                device_map="auto",
-                cache_dir=cache_dir,
-                torch_dtype=torch.float16,
-                is_trainable=False,
-            )
-    model.eval()
+    global model, tokenizer
+    if model is None:
+        from llama_rope_scaled_monkey_patch import replace_llama_rope_with_scaled_rope
+        replace_llama_rope_with_scaled_rope()
+        t_model = transformers.AutoModelForCausalLM.from_pretrained(
+                    base_model,
+                    torch_dtype=torch.float16,
+                    cache_dir=cache_dir,
+                    device_map="auto",
+                )
+        model = PeftModel.from_pretrained(
+                    t_model,
+                    lora_weights,
+                    device_map="auto",
+                    cache_dir=cache_dir,
+                    torch_dtype=torch.float16,
+                    is_trainable=False,
+                )
+        model.eval()
+        tokenizer =  AutoTokenizer.from_pretrained(base_model,use_fast=False,cache_dir=cache_dir)
+        tokenizer.pad_token = tokenizer.unk_token
+        model = model.to("cuda")
     return model
-    
+
+model, tokenizer = loadModel()
+
 #@spaces.GPU(duration=120)
 def respond(
     message,
@@ -96,10 +102,7 @@ def respond(
     max_tokens,
     temperature,
     top_p,
-):
-    model = loadModel()
-    tokenizer =  AutoTokenizer.from_pretrained(base_model,use_fast=False,cache_dir=cache_dir)
-    tokenizer.pad_token = tokenizer.unk_token
+):    
     ins_f = generate_prompt(message,None)
     inputs  =  tokenizer(ins_f, return_tensors="pt")
     input_ids = inputs["input_ids"].cuda()