Update app.py

app.py

@@ -301,8 +301,7 @@ demo.launch()'''
 
 
    
-
-'''import gradio as gr
+import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import sys
@@ -456,131 +455,4 @@ 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()'''
-
-      
-  
-      
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
-
-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} and provide two directions in which the story can go."
-    
-    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=400,
-        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()
-    
-    if "Option 1:" in response and "Option 2:" in response:
-        story_part = response.split("Option 1:")[0].strip()
-        options_part = "Option 1:" + response.split("Option 1:")[1]
-        option1 = options_part.split("Option 2:")[0].replace("Option 1:", "").strip()
-        option2 = options_part.split("Option 2:")[1].strip()
-    else:
-        story_part = response
-        option1 = "Follow the mysterious voice"
-        option2 = "Explore the abandoned house"
-    
-    return story_part, option1, option2
-
-def handle_option_click(current_story, option_text):
-    """Handle option selection and update story"""
-    if not option_text or not current_story:
-        return current_story, "Option 1", "Option 2"
-    
-    new_story_part, option1, option2 = generate_story_continuation(current_story, option_text)
-    updated_story = f"{current_story}\n\n⚡ You chose: {option_text}\n\n{new_story_part}"
-    return updated_story, option1, 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")
-        
-        story_display = gr.Textbox(
-            label="Your Story Journey",
-            lines=12,
-            interactive=False,
-            show_copy_button=True
-        )
-        
-    with gr.Row():
-        option1_btn = gr.Button("Option 1", visible=False)
-        option2_btn = gr.Button("Option 2", visible=False)
-    
-    exit_btn = gr.Button("Exit Story", variant="stop")
-    
-    def generate_initial_story(prompt):
-        story_part, option1, option2 = generate_story_continuation(prompt)
-        return {
-            story_display: story_part,
-            option1_btn: gr.update(value=option1, visible=True),
-            option2_btn: gr.update(value=option2, visible=True)
-        }
-    
-    generate_btn.click(
-        generate_initial_story,
-        inputs=story_prompt,
-        outputs=[story_display, option1_btn, option2_btn]
-    )
-    
-    def option_selected(option_text, current_story):
-        updated_story, new_opt1, new_opt2 = handle_option_click(current_story, option_text)
-        return {
-            story_display: updated_story,
-            option1_btn: gr.update(value=new_opt1),
-            option2_btn: gr.update(value=new_opt2)
-        }
-    
-    option1_btn.click(
-        fn=option_selected,
-        inputs=[option1_btn, story_display],
-        outputs=[story_display, option1_btn, option2_btn]
-    )
-    
-    option2_btn.click(
-        fn=option_selected,
-        inputs=[option2_btn, story_display],
-        outputs=[story_display, option1_btn, option2_btn]
-    )
-    
-    exit_btn.click(fn=lambda: None, inputs=None, outputs=None, js="window.close()")
-    
-    demo.launch()
+demo.launch()