Update app.py

app.py

@@ -301,7 +301,7 @@ demo.launch()'''
 
 
    
-import gradio as gr
+'''import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import sys
@@ -455,4 +455,128 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     # Exit handling
     exit_check.change(fn=lambda x: sys.exit() if x else None, inputs=[exit_check])
 
-demo.launch()
+demo.launch()'''
+
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import sys
+
+model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16,
+    device_map="auto",
+)
+
+def generate_story_continuation(current_story, selected_option=None):
+    """
+    Generate story continuation based on current story and selected option
+    """
+    if selected_option:
+        prompt = f"""Previous story: {current_story}
+        Selected path: {selected_option}
+        Continue the story with this choice and provide two new options."""
+    else:
+        prompt = f"Create a story based on: {current_story}"
+    
+    messages = [{"role": "user", "content": prompt}]
+    formatted_prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    
+    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(
+        **inputs,
+        max_length=200,
+        temperature=0.7,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id
+    )
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = response.split("Assistant: ")[-1].strip()
+    
+    try:
+        parts = response.split("Option 1:")
+        story = parts[0].strip()
+        options_text = "Option 1:" + parts[1]
+        option1 = options_text.split("Option 2:")[0].replace("Option 1:", "").strip()
+        option2 = options_text.split("Option 2:")[1].strip()
+    except:
+        story = response
+        option1 = "Path A: Continue the adventure"
+        option2 = "Path B: Take a different direction"
+    
+    return story, option1, option2
+
+def update_story(story, selected_option):
+    """
+    Update story based on selected option
+    """
+    new_story, new_option1, new_option2 = generate_story_continuation(story, selected_option)
+    full_story = f"{story}\n\nChosen path: {selected_option}\n\n{new_story}"
+    return full_story, new_option1, new_option2
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🌟 Interactive Story Generator")
+    
+    with gr.Column():
+        story_prompt = gr.Textbox(
+            label="Start your story",
+            placeholder="Enter your story prompt here...",
+            lines=3
+        )
+        
+        generate_btn = gr.Button(
+            "Generate Story",
+            variant="primary",
+            scale=1,
+            min_width=200,
+            size="lg"
+        )
+        
+        story_output = gr.Textbox(
+            label="Story",
+            lines=10,
+            show_copy_button=True
+        )
+        
+        with gr.Row():
+            option1_btn = gr.Button("Option 1")
+            option2_btn = gr.Button("Option 2")
+        
+        exit_check = gr.Checkbox(
+            label="Exit Story",
+            value=False
+        )
+    
+    def handle_option_selection(story, option_text):
+        return update_story(story, option_text)
+    
+    generate_btn.click(
+        fn=generate_story_continuation,
+        inputs=[story_prompt],
+        outputs=[story_output, option1_btn, option2_btn]
+    )
+    
+    option1_btn.click(
+        fn=handle_option_selection,
+        inputs=[story_output, option1_btn],
+        outputs=[story_output, option1_btn, option2_btn]
+    )
+    
+    option2_btn.click(
+        fn=handle_option_selection,
+        inputs=[story_output, option2_btn],
+        outputs=[story_output, option1_btn, option2_btn]
+    )
+    
+    exit_check.change(fn=lambda x: sys.exit() if x else None, inputs=[exit_check])
+
+demo.launch()
+
+