Update app.py

app.py

@@ -1,61 +1,70 @@
-# file: interactive_storytelling.py
 import streamlit as st
 from unsloth import FastLanguageModel
-from transformers import TextStreamer
-# import torch
+import torch
 
-# Initialize the model and tokenizer
+# Load the model and tokenizer
 max_seq_length = 2048
-dtype = None  # Auto-detection
-load_in_4bit = True
+dtype = None  # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
+load_in_4bit = True  # Use 4bit quantization to reduce memory usage. Can be False.
 
-# Load the pre-trained model
+# Initialize the model
 model, tokenizer = FastLanguageModel.from_pretrained(
     model_name="suhaif/unsloth-llama-3-8b-4bit",
     max_seq_length=max_seq_length,
     dtype=dtype,
-    load_in_4bit=load_in_4bit,
+    load_in_4bit=load_in_4bit
 )
 
-FastLanguageModel.for_inference(model)
+# Default instruction
+default_instruction = "You are a creative writer. Based on the given input, generate a well-structured story with an engaging plot, well-developed characters, and immersive details. Ensure the story has a clear beginning, middle, and end. Include dialogue and descriptions to bring the story to life. You can add a twist to the story also."
 
-# Set up the Streamlit app
+def format_prompt(input_text, instruction=default_instruction):
+    return f"{instruction}\n\nInput:\n{input_text}\n\nResponse:\n"
+
+# Streamlit App
 st.title("Interactive Storytelling Assistant")
-st.write("Collaboratively write stories and get real-time suggestions!")
-
-# Text input for the user prompt
-user_input = st.text_area("Write the beginning of your story:", placeholder="Once upon a time...")
-
-if st.button("Generate Suggestions"):
-    if user_input:
-        # Tokenize and send the input to the model
-        inputs = tokenizer([user_input], return_tensors="pt")
-
-        # Stream the generated output
-        text_streamer = TextStreamer(tokenizer)
-        st.write("Generating suggestions...")
-        generated_text = model.generate(**inputs, streamer=text_streamer, max_new_tokens=250)
-        
-        # Display the generated text
-        st.subheader("Story Suggestions:")
-        st.write(generated_text)
-
-        # Feedback section for rating
-        rating = st.slider("Rate the suggestion:", 1, 5)
-        if rating:
-            st.success(f"Thank you for your feedback! You rated this {rating} stars.")
-
-# Community sharing section
-st.subheader("Community Stories")
-st.write("Upload your story and share it with the community!")
+st.write("Create your story prompt and receive story suggestions!")
+
+# User input for story prompt
+user_input = st.text_area("Enter your story idea:", "A young adventurer embarks on a journey to find a lost treasure.")
+generate_story = st.button("Generate Story")
+
+if generate_story and user_input:
+    # Prepare inputs for the model
+    inputs = tokenizer(
+        [format_prompt(user_input)],
+        return_tensors="pt"
+    ).to("cuda")
+    
+    # Generate story
+    outputs = model.generate(**inputs, max_new_tokens=500, use_cache=True)
+    generated_story = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Display generated story
+    st.subheader("Generated Story:")
+    st.write(generated_story)
 
-uploaded_story = st.file_uploader("Upload your story (txt file)", type="txt")
+# Feedback mechanism
+st.subheader("Rate the Story")
+story_rating = st.slider("How would you rate this story?", 1, 5)
+user_feedback = st.text_area("Additional Feedback/Suggestions:")
 
-if uploaded_story:
-    # Display uploaded story
-    story_text = uploaded_story.read().decode("utf-8")
-    st.text_area("Your Uploaded Story:", story_text, height=300)
+if st.button("Submit Feedback"):
+    st.write("Thank you for your feedback!")
+    # Process feedback (In a real scenario, this would be saved to a database)
+
+# Community engagement
+st.subheader("Share Your Story")
+user_story = st.text_area("Write or paste your own story here:")
+if st.button("Share Story"):
+    st.write("Thank you for sharing your story!")
+    # Save the story (In a real scenario, this would be saved to a shared community platform)
+
+# Display shared stories (Placeholder example)
+st.subheader("Community Stories")
+st.write("Story 1: An epic tale of adventure...")  # Placeholder for actual stories
 
-# Display other community stories (could be expanded with database integration)
-st.subheader("Community Stories Feedback Section")
-st.write("In this section, you can provide critiques for other community stories.")
+# Critique section
+st.text_area("Leave a critique for Story 1:")
+if st.button("Submit Critique"):
+    st.write("Thank you for your critique!")