Update app.py

app.py

@@ -1,32 +1,22 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+from transformers import pipeline
 
 # Set up the Streamlit app
 st.title("🧚‍♀️ Magic Story Buddy 📚")
 st.markdown("Let's create a magical story just for you!")
 
-# Initialize the model
-@st.cache(allow_output_mutation=True)
-def load_model():
-    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)
+def initialize_model():
+    # Load the tokenizer and model
+    tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-2.7B")
+    model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-2.7B")
 
-        # Get configuration
-        config = AutoConfig.from_pretrained(model_name)
-        model_type = config.model_type
+    # Create the pipeline for text generation
+    text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-        print(f"Loaded model type: {model_type}")
+    return text_generator
 
-        return model, tokenizer
-
-    except Exception as e:
-        st.error(f"Error loading model: {e}")
-
-model, tokenizer = load_model()
+# Initialize the model and pipeline
+pipe = initialize_model()
 
 # User input
 child_name = st.text_input("What's your name, young storyteller?")
@@ -38,9 +28,9 @@ 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)
+    messages = [{"role": "user", "content": prompt}]
+    result = pipe(messages, max_length=max_length, do_sample=True, temperature=0.7)
+    return result[0]['generated_text']
 
 if st.button("Create My Story!"):
     if child_name and story_theme:
@@ -51,7 +41,6 @@ if st.button("Create My Story!"):
         - 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}. """