Update app.py

app.py

@@ -1,22 +1,13 @@
 import streamlit as st
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+import openai
+import os
 
 # Set up the Streamlit app
 st.title("🧚‍♀️ Magic Story Buddy 📚")
 st.markdown("Let's create a magical story just for you!")
 
-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")
-
-    # Create the pipeline for text generation
-    text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-    return text_generator
-
-# Initialize the model and pipeline
-pipe = pipeline("text-generation", model="hakurei/lit-6B")
+# Set up OpenAI API key
+openai.api_key = os.getenv("sk-proj-tWGLTJozNuRBlh6RGqFAT3BlbkFJQy7dSGpukXfwuwueZGIt")
 
 # User input
 child_name = st.text_input("What's your name, young storyteller?")
@@ -27,10 +18,19 @@ 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):
-    messages = [{"role": "user", "content": prompt}]
-    result = pipe(messages, max_length=max_length, do_sample=True, temperature=0.7)
-    return result[0]['generated_text']
+def generate_story(prompt, max_tokens=500):
+    response = openai.ChatCompletion.create(
+        model="gpt-3.5-turbo",
+        messages=[
+            {"role": "system", "content": "You are a friendly storyteller for children."},
+            {"role": "user", "content": prompt}
+        ],
+        max_tokens=max_tokens,
+        n=1,
+        stop=None,
+        temperature=0.7,
+    )
+    return response.choices[0].message['content'].strip()
 
 if st.button("Create My Story!"):
     if child_name and story_theme:
@@ -48,7 +48,7 @@ if st.button("Create My Story!"):
             prompt += "This story teaches us that "
         
         # Generate the story
-        story = generate_story(prompt, max_length=story_length)
+        story = generate_story(prompt, max_tokens=story_length * 2)  # Multiply by 2 as tokens != words
         
         # Display the story
         st.markdown("## Your Magical Story")