Update app.py

app.py

@@ -18,8 +18,8 @@ DESCRIPTION = """
 """
 
 #tokenizer = AutoTokenizer.from_pretrained(model_id, device_map="auto", trust_remote_code=True)
-tokenizer = LlamaTokenizerFast.from_pretrained(tokenizer_id, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", trust_remote_code=True)
+tokenizer = LlamaTokenizerFast.from_pretrained(tokenizer_id, device_map="cuda", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", trust_remote_code=True)
 
 def format_prompt(user_message, system_message="You are Grok-1, an AI language model created by Tonic-AI. You are a cautious assistant. You carefully follow instructions. You are helpful and harmless and follow ethical guidelines and promote positive behavior.\n\n"):
 #   prompt = f"<|im_start|>assistant\n{system_message}<|im_end|>\n<|im_start|>\nuser\n{user_message}<|im_end|>\nassistant\n"