Update app.py

app.py

@@ -114,6 +114,33 @@ def load_llm_model(model_id, hf_token):
         print("Error: HF_TOKEN secret is not set. Cannot load Hugging Face model.")
         return None, None
 
+    try:
+        llm_tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
+
+        # Explicitly set the chat template for Gemma models
+        # This template formats messages as <start_of_turn>user\n{message}<end_of_turn>\n<start_of_turn>model\n{response}<end_of_turn>\n
+        # and adds <bos> at the beginning and <start_of_turn>model\n at the end for generation prompt.
+        llm_tokenizer.chat_template = "{% for message in messages %}{% if message['role'] == 'user' %}{{ '<start_of_turn>user\n' + message['content'] + '<end_of_turn>\n' }}{% elif message['role'] == 'system' %}{{ '<start_of_turn>system\n' + message['content'] + '<end_of_turn>\n' }}{% elif message['role'] == 'tool' %}{{ '<start_of_turn>tool\n' + message['content'] + '<end_of_turn>\n' }}{% elif message['role'] == 'model' %}{{ '<start_of_turn>model\n' + message['content'] + '<end_of_turn>\n' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<start_of_turn>model\n' }}{% endif %}"
+
+
+        if llm_tokenizer.pad_token is None:
+             llm_tokenizer.pad_token = llm_tokenizer.eos_token
+
+        llm_model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            token=hf_token,
+            device_map="auto", # This will likely map to 'cpu'
+        )
+
+        print(f"Model {model_id} loaded in full precision.")
+        return llm_model, llm_tokenizer
+
+    except Exception as e:
+        print(f"Error loading model {model_id}: {e}")
+        print("Please ensure transformers, trl, peft, and accelerate are installed.")
+        print("Check your Hugging Face token.")
+        return None, None
+
     try:
         llm_tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
         if llm_tokenizer.pad_token is None: