Update app.py

app.py

@@ -1,20 +1,32 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
 import torch
 
 def load_model():
-    # Fixed the spelling in the model name
-    model_name = "GouthamVarma/mentalhealth_coversational_chatbot"  
+    model_name = "GouthamVarma/mentalhealth_coversational_chatbot"
+    base_model = "google/gemma-2b-it"  # Base model we fine-tuned from
     
-    # Load model without authentication since it's public
-    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    # First load the base model's tokenizer and config
+    tokenizer = AutoTokenizer.from_pretrained(
+        base_model, 
+        trust_remote_code=True
+    )
+    
+    config = AutoConfig.from_pretrained(
+        base_model,
+        trust_remote_code=True
+    )
+    
+    # Then load your fine-tuned model
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
+        config=config,
         low_cpu_mem_usage=True,
         torch_dtype=torch.float32,
         device_map="cpu",
         trust_remote_code=True
     )
+    
     return model, tokenizer
 
 print("Loading model... This might take a few minutes...")