Update app.py

app.py

@@ -27,10 +27,18 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # model_id = "mistralai/Mistral-7B-v0.3"
 
-model_id = "CohereForAI/c4ai-command-r-plus-4bit"
+model_id = "CohereForAI/c4ai-command-r-plus"
 
 
+tokenizer = AutoTokenizer.from_pretrained(model_id, token= token)
 
+model = AutoModelForCausalLM.from_pretrained(model_id, token= token, 
+                                                 torch_dtype=torch.bfloat16,    
+                                             attn_implementation="flash_attention_2",
+                                             # low_cpu_mem_usage=True,
+                                             # llm_int8_enable_fp32_cpu_offload=True,
+                                             device_map="cuda"
+                                            )
 
 
 # 
@@ -46,15 +54,7 @@ def respond(
     temperature,
     top_p,
 ):
-    tokenizer = AutoTokenizer.from_pretrained(model_id, token= token)
 
-    model = AutoModelForCausalLM.from_pretrained(model_id, token= token, 
-                                                 # torch_dtype=torch.bfloat16,    
-                                             # attn_implementation="flash_attention_2",
-                                             # low_cpu_mem_usage=True,
-                                             # llm_int8_enable_fp32_cpu_offload=True,
-                                             device_map="cuda"
-                                            )
     messages = [{"role": "user", "content": "Hello, how are you?"}]
     input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to('cuda')
 ## <BOS_TOKEN><|START_OF_TURN_TOKEN|><|USER_TOKEN|>Hello, how are you?<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>