Update app.py

app.py

@@ -12,11 +12,11 @@ MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
-
-model_id = "stabilityai/ar-stablelm-2-chat"
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-tokenizer.use_default_system_prompt = False
+def load_model():
+    model_id = "stabilityai/ar-stablelm-2-chat"
+    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    tokenizer.use_default_system_prompt = False
 
 
 def generate(
@@ -113,6 +113,7 @@ with gr.Blocks(css_paths="style.css", fill_height=True) as demo:
         try:
             # Validate token using Hugging Face Hub API
             login(token = hf_token)
+            load_model()
             return f"Authenticated successfully"
         except HfHubHTTPError:
             return "Invalid token. Please try again."