Update app.py

app.py

@@ -1,15 +1,19 @@
 import gradio as gr
-from transformers import pipeline, set_seed
+from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed
+import torch
 
-# Load a safer, creative chat model
-generator = pipeline(
-    "text-generation",
-    model="Qwen/Qwen1.5-1.8B-Chat",
-    device=-1
-)
+MODEL_NAME = "Qwen/Qwen1.5-1.8B-Chat"
 
 set_seed(42)
 
+# Load tokenizer & model correctly
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float32,
+    device_map="auto"
+)
+
 def get_actions(theme):
     return {
         "Fantasy": [
@@ -29,66 +33,86 @@ def get_actions(theme):
         ],
     }[theme]
 
-def generate_story(prompt):
-    output = generator(
-        prompt,
-        max_new_tokens=350,       # LONGER story
-        temperature=0.95,         # creative but safe
+def generate_story(messages):
+    # Apply proper chat template
+    input_ids = tokenizer.apply_chat_template(
+        messages,
+        return_tensors="pt"
+    ).to(model.device)
+
+    output_ids = model.generate(
+        input_ids,
+        max_new_tokens=400,        # LONG story
+        temperature=0.95,          # creative but calm
         top_p=0.9,
-        repetition_penalty=1.12,
+        repetition_penalty=1.1,
         do_sample=True,
-        eos_token_id=generator.tokenizer.eos_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+
+    output = tokenizer.decode(
+        output_ids[0][input_ids.shape[-1]:],
+        skip_special_tokens=True
     )
-    return output[0]["generated_text"].strip()
+
+    return output.strip()
 
 def start_story(name, theme):
-    prompt = f"""
-You are a gentle, imaginative storyteller.
-
-STRICT RULES:
-- Write ONLY the story
-- Do NOT ask questions
-- Do NOT ask for feedback
-- No violence, threats, or fear
-- Keep the tone cozy, adventurous, and magical
-- Stay strictly in the {theme} genre
-
-STORY:
-Begin a vivid and creative {theme} story.
-The main character is named {name}.
-Use rich descriptions, warm emotions, and wonder.
-"""
-
-    story = generate_story(prompt)
+    messages = [
+        {
+            "role": "system",
+            "content": (
+                "You are a gentle, imaginative storyteller.\n"
+                "Write ONLY the story.\n"
+                "Never ask questions.\n"
+                "Never request feedback.\n"
+                "No violence, threats, or dark content.\n"
+                "Keep the tone cozy, magical, and adventurous."
+            )
+        },
+        {
+            "role": "user",
+            "content": (
+                f"Begin a vivid {theme} story.\n"
+                f"The main character is named {name}.\n"
+                "Use rich descriptions and warm emotions."
+            )
+        }
+    ]
+
+    story = generate_story(messages)
     return story, gr.update(choices=get_actions(theme), visible=True), ""
 
 def next_step(name, theme, action):
     if not action:
         return "Please select an action to continue the story.", ""
 
-    prompt = f"""
-You are continuing a story.
-
-STRICT RULES:
-- ONLY continue the story
-- No questions or feedback requests
-- No violence, threats, or dark themes
-- Keep it imaginative, light, and engaging
-- Stay in the {theme} genre
-
-STORY CONTEXT:
-The main character is {name}.
-{name} chooses to {action}.
-
-Continue the story with rich detail and a sense of wonder.
-"""
-
-    story = generate_story(prompt)
+    messages = [
+        {
+            "role": "system",
+            "content": (
+                "You are continuing a cozy, imaginative story.\n"
+                "Write ONLY the story.\n"
+                "No violence or threatening language."
+            )
+        },
+        {
+            "role": "user",
+            "content": (
+                f"The story is set in a {theme} world.\n"
+                f"The main character is {name}.\n"
+                f"{name} chooses to {action}.\n"
+                "Continue the story with wonder and detail."
+            )
+        }
+    ]
+
+    story = generate_story(messages)
     return story, ""
 
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo")) as demo:
-    gr.Markdown("# 📖 Cozy Infinite AI Story Generator")
-    gr.Markdown("Creative, gentle, and endlessly imaginative stories ✨")
+    gr.Markdown("# 📖 Cozy Creative Story Generator")
+    gr.Markdown("Safe, imaginative, long-form storytelling ✨")
 
     with gr.Row():
         with gr.Column():
@@ -97,7 +121,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo")) as demo:
             start_btn = gr.Button("🚀 Start Story")
 
         with gr.Column():
-            story_out = gr.Textbox(label="Story", lines=12, interactive=False)
+            story_out = gr.Textbox(label="Story", lines=14, interactive=False)
             action_choice = gr.Dropdown(label="Your Action", choices=[], visible=False)
             continue_btn = gr.Button("✨ Continue Story")