Update app.py

app.py

@@ -584,158 +584,206 @@ demo.launch()'''
    
    
 
-
- 
-
-  
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import random
 
 # Load TinyLlama model and tokenizer
 model_name = "TinyLlama/TinyLlama-1.1B-Chat"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
-
-# Initialize story state
-story_history = []
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
 
-def generate_text(prompt, max_length=100):
-    """Generate text using TinyLlama model with optimized settings."""
-    # Format prompt for chat
-    formatted_prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
-    
-    # Tokenize input
-    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
-    
-    # Generate with optimized parameters
-    with torch.no_grad():
-        outputs = model.generate(
-            inputs.input_ids,
-            max_new_tokens=max_length,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9,
-            num_return_sequences=1,
-            pad_token_id=tokenizer.eos_token_id
-        )
-    
-    # Decode and clean up response
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    response = response.replace(formatted_prompt, "").strip()
-    return response
+class StoryState:
+    def __init__(self):
+        self.history = []
+        self.player_name = ""
+        self.character_class = ""
+        self.current_location = ""
+        self.inventory = []
+        self.health = 100
+        
+    def reset(self):
+        self.__init__()
 
-def generate_initial_story(prompt):
-    """Generate a coherent initial story paragraph from the prompt."""
-    story_prompt = (
-        f"Write a short, engaging paragraph (40-50 words) for a story starting with: '{prompt}'. "
-        "Make it descriptive and ensure it relates directly to the prompt."
-    )
-    story_output = generate_text(story_prompt, max_length=75)
-    
-    # Ensure the story starts with the prompt theme
-    if not any(word in story_output.lower() for word in prompt.lower().split()):
-        story_output = f"{prompt}. {story_output}"
-    
-    return story_output
+# Initialize game state
+game_state = StoryState()
 
-def generate_options(story_so_far):
-    """Generate three context-aware options based on the current story."""
-    option_prompt = (
-        f"Based on this story: '{story_so_far}'\n"
-        "Generate 3 distinct, exciting options for what happens next. "
-        "Each option should be a single sentence and directly relate to the story's context. "
-        "Format: 1. [Option 1] 2. [Option 2] 3. [Option 3]"
-    )
-    
-    options_text = generate_text(option_prompt, max_length=100)
-    
-    # Parse options or provide fallback options if generation fails
+def generate_story_segment(prompt, max_length=150):
+    """Generate story text using TinyLlama."""
     try:
-        options = [opt.split(". ")[1].strip() for opt in options_text.split("\n") if opt.strip().startswith(("1.", "2.", "3."))]
-        if len(options) != 3:
-            raise ValueError("Invalid options generated")
-    except:
-        # Fallback options based on common story elements
-        character = story_so_far.split()[0]
-        options = [
-            f"{character} faces an unexpected challenge that tests their resolve.",
-            f"{character} discovers something that changes everything.",
-            f"{character} makes a decision that leads to adventure."
+        # Format prompt for chat
+        formatted_prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
+        
+        # Tokenize
+        inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
+        
+        # Generate with optimized parameters
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs.input_ids,
+                max_new_tokens=max_length,
+                do_sample=True,
+                temperature=0.8,
+                top_p=0.9,
+                num_return_sequences=1,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        # Decode and clean response
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response.replace(formatted_prompt, "").strip()
+    except Exception as e:
+        return get_fallback_segment()
+
+def get_fallback_segment():
+    """Provide fallback story segments if generation fails."""
+    fallbacks = [
+        "As you proceed cautiously, the path ahead seems to hold both promise and peril.",
+        "The atmosphere grows tense as you consider your next move.",
+        "Your instincts tell you that this decision could change everything.",
+        "The consequences of your choices begin to unfold around you."
+    ]
+    return random.choice(fallbacks)
+
+def generate_choices(context):
+    """Generate contextual choices based on the current story state."""
+    try:
+        prompt = (
+            f"Based on this story context: '{context}'\n"
+            "Generate 3 distinct, interesting choices for what the player could do next. "
+            "Each choice should be a single sentence and affect the story differently. "
+            "Format: 1. [Choice 1] 2. [Choice 2] 3. [Choice 3]"
+        )
+        
+        choices_text = generate_story_segment(prompt, max_length=100)
+        choices = [
+            choice.split(". ")[1].strip() 
+            for choice in choices_text.split("\n") 
+            if choice.strip().startswith(("1.", "2.", "3."))
         ]
+        
+        if len(choices) != 3:
+            raise ValueError("Invalid choices generated")
+            
+        return choices
+    except:
+        return get_fallback_choices()
+
+def get_fallback_choices():
+    """Provide fallback choices if generation fails."""
+    return [
+        "Proceed carefully and gather more information",
+        "Take bold action and face the consequences",
+        "Seek an alternative path forward"
+    ]
+
+def initialize_character(name, character_class):
+    """Initialize the player character."""
+    if not name or not character_class:
+        return "Please enter both name and class!", None, None, None
+    
+    game_state.reset()
+    game_state.player_name = name
+    game_state.character_class = character_class
+    
+    # Generate initial story setup
+    prompt = (
+        f"Write an engaging opening paragraph for a fantasy adventure story where {name}, "
+        f"a brave {character_class}, begins their journey. Make it exciting and descriptive."
+    )
     
-    return options
-
-def start_story(initial_prompt):
-    """Start the story with the user's initial prompt."""
-    global story_history
-    story_history = [initial_prompt]
-    
-    # Generate initial segment
-    segment = generate_initial_story(initial_prompt)
-    story_history.append(segment)
+    story_start = generate_story_segment(prompt)
+    game_state.history.append(story_start)
     
-    # Generate options
-    options = generate_options(segment)
+    # Generate initial choices
+    choices = generate_choices(story_start)
     
-    return format_story_output(segment, options), *options
+    return format_game_output(story_start, choices), *choices
 
 def continue_story(choice):
-    """Continue the story based on the user's choice."""
-    global story_history
-    story_history.append(f"Choice: {choice}")
-    
-    # Generate new segment
-    continuation_prompt = (
-        f"Continue this story coherently:\n{story_history[-2]}\n"
-        f"The chosen path: {choice}\n"
-        "Write a natural 40-50 word continuation that builds on these events."
+    """Continue the story based on player's choice."""
+    if not choice or not game_state.history:
+        return "Please start a new game!", None, None, None
+    
+    game_state.history.append(f"You chose to: {choice}")
+    
+    # Generate new story segment
+    prompt = (
+        f"Continue this story naturally:\n{game_state.history[-2]}\n"
+        f"The player ({game_state.player_name} the {game_state.character_class}) chose to: {choice}\n"
+        "Write an exciting 50-75 word continuation that builds on this choice."
     )
     
-    new_segment = generate_text(continuation_prompt, max_length=75)
-    story_history.append(new_segment)
+    new_segment = generate_story_segment(prompt)
+    game_state.history.append(new_segment)
     
-    # Generate new options
-    options = generate_options(new_segment)
+    # Generate new choices
+    choices = generate_choices(new_segment)
     
-    return format_story_output("\n".join(story_history), options), *options
+    return format_game_output("\n\n".join(game_state.history), choices), *choices
 
-def format_story_output(story, options):
-    """Format the story and options for display."""
-    return f"{story}\n\n🔹 What happens next?\n1️⃣ {options[0]}\n2️⃣ {options[1]}\n3️⃣ {options[2]}"
+def format_game_output(story, choices):
+    """Format the game output with story and choices."""
+    status_line = (
+        f"📜 Adventure of {game_state.player_name} the {game_state.character_class}\n"
+        f"❤️ Health: {game_state.health}% | 🎒 Items: {', '.join(game_state.inventory) or 'None'}\n"
+        f"{'─' * 50}\n\n"
+    )
+    
+    return (
+        f"{status_line}{story}\n\n"
+        f"What will you do?\n"
+        f"1️⃣ {choices[0]}\n"
+        f"2️⃣ {choices[1]}\n"
+        f"3️⃣ {choices[2]}"
+    )
 
-def reset_story():
-    """Reset the story to start fresh."""
-    global story_history
-    story_history = []
-    return "Story reset. Enter a new prompt to begin!", "", "", ""
+def reset_game():
+    """Reset the game state."""
+    game_state.reset()
+    return "Start a new adventure!", "", "", "", None, None, None
 
 # Gradio interface
-with gr.Blocks(title="Story Adventure") as demo:
-    gr.Markdown("# 📚 Story Adventure")
-    gr.Markdown("Create your own adventure! Enter a prompt and shape the story through your choices.")
+with gr.Blocks(title="Story Quest") as demo:
+    gr.Markdown("# 🎮 Story Quest")
+    gr.Markdown("Embark on an AI-powered interactive adventure! Create your character and shape the story through your choices.")
     
+    # Character Creation
     with gr.Row():
-        prompt_input = gr.Textbox(
-            label="Story Prompt",
-            placeholder="e.g., 'Sarah discovers a mysterious letter in her attic'",
-            lines=2
-        )
-        start_button = gr.Button("🎮 Start Story", variant="primary")
+        with gr.Column():
+            name_input = gr.Textbox(
+                label="Character Name",
+                placeholder="Enter your character's name"
+            )
+            class_input = gr.Dropdown(
+                choices=["Warrior", "Mage", "Rogue", "Paladin", "Ranger"],
+                label="Character Class"
+            )
+            start_button = gr.Button("⚔️ Begin Adventure", variant="primary")
     
-    story_output = gr.Textbox(label="Your Story", lines=12, interactive=False)
+    # Story Display and Choices
+    story_output = gr.Textbox(
+        label="Your Adventure",
+        lines=12,
+        interactive=False
+    )
     
     with gr.Row():
-        choice_button_1 = gr.Button(value="", visible=False, variant="secondary")
-        choice_button_2 = gr.Button(value="", visible=False, variant="secondary")
-        choice_button_3 = gr.Button(value="", visible=False, variant="secondary")
+        choice_button_1 = gr.Button(visible=False, variant="secondary")
+        choice_button_2 = gr.Button(visible=False, variant="secondary")
+        choice_button_3 = gr.Button(visible=False, variant="secondary")
     
-    reset_button = gr.Button("🔄 Reset Story", variant="stop")
+    reset_button = gr.Button("🔄 Reset Adventure", variant="stop")
     
     # Event handlers
     start_button.click(
-        fn=start_story,
-        inputs=prompt_input,
+        fn=initialize_character,
+        inputs=[name_input, class_input],
         outputs=[story_output, choice_button_1, choice_button_2, choice_button_3]
     ).then(
         fn=lambda story, c1, c2, c3: (
@@ -751,7 +799,7 @@ with gr.Blocks(title="Story Adventure") as demo:
     for button in [choice_button_1, choice_button_2, choice_button_3]:
         button.click(
             fn=continue_story,
-            inputs=button,
+            inputs=[button],
             outputs=[story_output, choice_button_1, choice_button_2, choice_button_3]
         ).then(
             fn=lambda story, c1, c2, c3: (
@@ -765,9 +813,23 @@ with gr.Blocks(title="Story Adventure") as demo:
         )
     
     reset_button.click(
-        fn=reset_story,
-        outputs=[story_output, choice_button_1, choice_button_2, choice_button_3]
+        fn=reset_game,
+        outputs=[
+            story_output,
+            name_input,
+            class_input,
+            story_output,
+            choice_button_1,
+            choice_button_2,
+            choice_button_3
+        ]
     )
 
+# Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()
+ 
+
+  
+
+