Updated mistral model

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 # Set up the Streamlit app
 st.title("🧚‍♀️ Magic Story Buddy 📚")
@@ -8,9 +9,11 @@ st.markdown("Let's create a magical story just for you!")
 # Initialize the model
 @st.cache_resource
 def load_model():
-    return pipeline("text-generation", model="ajibawa-2023/Young-Children-Storyteller-Mistral-7B")
+    model = AutoModelForCausalLM.from_pretrained("ajibawa-2023/Young-Children-Storyteller-Mistral-7B", torch_dtype=torch.float16)
+    tokenizer = AutoTokenizer.from_pretrained("ajibawa-2023/Young-Children-Storyteller-Mistral-7B")
+    return model, tokenizer
 
-model = load_model()
+model, tokenizer = load_model()
 
 # User input
 child_name = st.text_input("What's your name, young storyteller?")
@@ -21,15 +24,29 @@ story_theme = st.selectbox("What would you like your story to be about?",
 story_length = st.slider("How long should the story be?", 50, 200, 100)
 include_moral = st.checkbox("Include a moral lesson?")
 
+def generate_story(prompt, max_length=500):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(**inputs, max_length=max_length, num_return_sequences=1, do_sample=True, temperature=0.7)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
 if st.button("Create My Story!"):
     if child_name and story_theme:
         # Construct the prompt
-        prompt = f"[CHILDREN'S STORY] Once upon a time, in a {story_theme.lower()}, there was a brave child named {child_name}. "
+        prompt = f"""Create a short children's story with the following details:
+        - Main character: {child_name}
+        - Theme: {story_theme}
+        - Length: About {story_length} words
+        - Audience: Children aged 5-10
+        - Tone: Friendly, educational, and imaginative
+
+        Story:
+        Once upon a time, in a {story_theme.lower()}, there was a brave child named {child_name}. """
+        
         if include_moral:
             prompt += "This story teaches us that "
         
         # Generate the story
-        story = model(prompt, max_length=story_length, num_return_sequences=1)[0]['generated_text']
+        story = generate_story(prompt, max_length=story_length)
         
         # Display the story
         st.markdown("## Your Magical Story")