Update app final mistral

app.py

@@ -7,11 +7,24 @@ st.title("🧚‍♀️ Magic Story Buddy 📚")
 st.markdown("Let's create a magical story just for you!")
 
 # Initialize the model
-@st.cache_resource
+@st.cache(allow_output_mutation=True)
 def load_model():
-    model = AutoModelForCausalLM.from_pretrained("blockblockblock/Young-Children-Storyteller-Mistral-7B-bpw6", torch_dtype=torch.float16)
-    tokenizer = AutoTokenizer.from_pretrained("blockblockblock/Young-Children-Storyteller-Mistral-7B-bpw6")
-    return model, tokenizer
+    model_name = "blockblockblock/Young-Children-Storyteller-Mistral-7B-bpw6"
+    try:
+        # Load model and tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+
+        # Get configuration
+        config = AutoConfig.from_pretrained(model_name)
+        model_type = config.model_type
+
+        print(f"Loaded model type: {model_type}")
+
+        return model, tokenizer
+
+    except Exception as e:
+        st.error(f"Error loading model: {e}")
 
 model, tokenizer = load_model()