app.py

import gradio as gr
import random
from datetime import datetime
import tempfile
import os
import edge_tts
import asyncio
import warnings
import pytz
import re
import json
import pandas as pd
from pathlib import Path
from gradio_client import Client

warnings.filterwarnings('ignore')

# Initialize story starters from GPT 
# Initialize story starters from Claude 

STORY_STARTERS = [
    
    ["Adventure", "A lost map leads a group of unlikely friends on a treasure hunt through ancient ruins."],
    ["Mystery", "A voice from an old radio transmits cryptic messages every night at 2 AM."],
    ["Romance", "An artist begins receiving love letters from a stranger who knows every detail about their paintings."],
    ["Sci-Fi", "In a future where memories can be traded, one person uncovers a memory that shouldn’t exist."],
    ["Fantasy", "A seemingly ordinary child finds out they have the power to heal people but at a mysterious cost."],
    ["Comedy-Sitcom", "The new roommate is convinced they are a spy, despite their day job at a coffee shop."],
    ["Comedy-Workplace", "A filing cabinet at the office is actually a portal to another dimension."],
    ["Comedy-Family", "A grandmother enters a contest to out-tech her tech-savvy grandkids and win the latest gadget."],
    ["Comedy-Supernatural", "A vampire tries to survive in a world obsessed with healthy eating and organic lifestyles."],
    ["Comedy-Travel", "An app glitch reroutes a couple on vacation to the world's quirkiest landmarks."],

    ['Historical Epic', 'In 1922, when Howard Carter discovered King Tut\'s tomb, he also found a sealed chamber with strange inscriptions that nobody could translate...'],
    ['Science Thriller', 'A quantum physicist working on parallel universe theory suddenly starts receiving messages from her alternate selves, all warning of an imminent catastrophe...'],
    ['True Crime Drama', 'While renovating the infamous Winchester Mystery House, workers discover a previously unknown room containing a journal dated 2025...'],
    ['Space Adventure', 'Based on the real Wow! Signal of 1977, a modern radio astronomer detects the same 72-second sequence repeating from a different location in space...'],
    ['Psychological Mystery', 'A renowned memory researcher developing a treatment for Alzheimer\'s realizes her patients are all remembering the same event that hasn\'t happened yet...'],
    ['Environmental Thriller', 'Deep in the Arctic permafrost, scientists studying climate change uncover a perfectly preserved prehistoric creature that begins to thaw...'],
    ['Tech Noir', 'When a pioneering AI researcher disappears, the only clue is a series of impossibly old photographs showing her throughout history...'],
    ['Maritime Mystery', 'Inspired by the Bermuda Triangle incidents, a marine archaeologist finds a modern submarine in a 200-year-old shipwreck...'],
    ['Cultural Adventure', 'While documenting the uncontacted tribes of the Amazon, an anthropologist discovers ancient cave paintings depicting modern technology...'],
    ['Medical Drama', 'During a routine genome sequencing study, a geneticist identifies a DNA pattern that seems to predict exactly when people will have life-changing moments...']

]

# Initialize story starters with added comedy section
STORY_STARTERS_ORIGINAL_RECIPE = [
    ['Adventure', 'In a hidden temple deep in the Amazon...'],
    ['Mystery', 'The detective found an unusual note...'],
    ['Romance', 'Two strangers meet on a rainy evening...'],
    ['Sci-Fi', 'The space station received an unexpected signal...'],
    ['Fantasy', 'A magical portal appeared in the garden...'],
    ['Comedy-Sitcom', 'The new roommate arrived with seven emotional support animals...'],
    ['Comedy-Workplace', 'The office printer started sending mysterious messages...'],
    ['Comedy-Family', 'Grandma decided to become a social media influencer...'],
    ['Comedy-Supernatural', 'The ghost haunting the house was absolutely terrible at scaring people...'],
    ['Comedy-Travel', 'The GPS insisted on giving directions in interpretive dance descriptions...']
]

# Initialize client outside of interface definition
arxiv_client = None

def init_client():
    global arxiv_client
    if arxiv_client is None:
        arxiv_client = Client("awacke1/Arxiv-Paper-Search-And-QA-RAG-Pattern")
    return arxiv_client

def save_story(story, audio_path):
    """Save story and audio to gallery with markdown formatting"""
    try:
        # Create gallery directory if it doesn't exist
        gallery_dir = Path("gallery")
        gallery_dir.mkdir(exist_ok=True)
        
        # Generate timestamp and sanitize first line for filename
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        first_line = story.split('\n')[0].strip()
        safe_name = re.sub(r'[^\w\s-]', '', first_line)[:50]  # First 50 chars, sanitized
        
        # Save story text as markdown
        story_path = gallery_dir / f"story_{timestamp}_{safe_name}.md"
        with open(story_path, "w") as f:
            f.write(f"# {first_line}\n\n{story}")
        
        # Copy audio file to gallery with matching name
        new_audio_path = None
        if audio_path:
            new_audio_path = gallery_dir / f"audio_{timestamp}_{safe_name}.mp3"
            os.system(f"cp {audio_path} {str(new_audio_path)}")
            
        return str(story_path), str(new_audio_path) if new_audio_path else None
    except Exception as e:
        print(f"Error saving to gallery: {str(e)}")
        return None, None

def load_gallery():
    """Load all stories and audio from gallery with markdown support"""
    try:
        gallery_dir = Path("gallery")
        if not gallery_dir.exists():
            return []
        
        files = []
        for story_file in sorted(gallery_dir.glob("story_*.md"), reverse=True):
            # Extract timestamp and name from filename
            parts = story_file.stem.split('_', 2)
            timestamp = f"{parts[1]}"
            
            # Find matching audio file
            audio_pattern = f"audio_{timestamp}_*.mp3"
            audio_files = list(gallery_dir.glob(audio_pattern))
            audio_file = audio_files[0] if audio_files else None
            
            # Read story content and get preview
            with open(story_file) as f:
                content = f.read()
                # Skip markdown header and get preview
                preview = content.split('\n\n', 1)[1][:100] + "..."
            
            files.append([
                timestamp,
                f"[{preview}]({str(story_file)})",  # Markdown link to story
                str(story_file),
                str(audio_file) if audio_file else None
            ])
        
        return files
    except Exception as e:
        print(f"Error loading gallery: {str(e)}")
        return []

# Keep all other functions unchanged
def generate_story(prompt, model_choice):
    """Generate story using specified model"""
    try:
        client = init_client()
        if client is None:
            return "Error: Story generation service is not available."
        
        result = client.predict(
            prompt=prompt,
            llm_model_picked=model_choice,
            stream_outputs=True,
            api_name="/ask_llm"
        )
        return result
    except Exception as e:
        return f"Error generating story: {str(e)}"

async def generate_speech(text, voice="en-US-AriaNeural"):
    """Generate speech from text"""
    try:
        communicate = edge_tts.Communicate(text, voice)
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp_file:
            tmp_path = tmp_file.name
            await communicate.save(tmp_path)
        return tmp_path
    except Exception as e:
        print(f"Error in text2speech: {str(e)}")
        return None

def process_story_and_audio(prompt, model_choice):
    """Process story, generate audio, and save to gallery"""
    try:
        # Generate story
        story = generate_story(prompt, model_choice)
        if isinstance(story, str) and story.startswith("Error"):
            return story, None, None
        
        # Generate audio
        audio_path = asyncio.run(generate_speech(story))
        
        # Save to gallery
        story_path, saved_audio_path = save_story(story, audio_path)
        
        return story, audio_path, load_gallery()
    except Exception as e:
        return f"Error: {str(e)}", None, None

def play_gallery_audio(evt: gr.SelectData, gallery_data):
    """Play audio from gallery selection"""
    try:
        selected_row = gallery_data[evt.index[0]]
        audio_path = selected_row[3]  # Audio path is the fourth element
        if audio_path and os.path.exists(audio_path):
            return audio_path
        return None
    except Exception as e:
        print(f"Error playing gallery audio: {str(e)}")
        return None

# Create the Gradio interface (keep unchanged)
with gr.Blocks(title="AI Story Generator") as demo:
    gr.Markdown("""
    # 🎭 AI Story Generator & Narrator
    Generate creative stories, listen to them, and build your gallery!
    """)
    
    with gr.Row():
        with gr.Column(scale=3):
            with gr.Row():
                prompt_input = gr.Textbox(
                    label="Story Concept",
                    placeholder="Enter your story idea...",
                    lines=3
                )
            
            with gr.Row():
                model_choice = gr.Dropdown(
                    label="Model",
                    choices=[
                        "mistralai/Mixtral-8x7B-Instruct-v0.1",
                        "mistralai/Mistral-7B-Instruct-v0.2"
                    ],
                    value="mistralai/Mixtral-8x7B-Instruct-v0.1"
                )
                generate_btn = gr.Button("Generate Story")
            
            with gr.Row():
                story_output = gr.Textbox(
                    label="Generated Story",
                    lines=10,
                    interactive=False
                )
            
            with gr.Row():
                audio_output = gr.Audio(
                    label="Story Narration",
                    type="filepath"
                )
        
        # Sidebar with Story Starters and Gallery
        with gr.Column(scale=1):
            gr.Markdown("### 📚 Story Starters")
            
            gr.Markdown("# 🎯 𝓜𝓲𝔁𝓽𝓻𝓪𝓵 𝓸𝓯 𝓔𝔁𝓹𝓮𝓻𝓽𝓼 ⚡")
            gr.Markdown("**Abstract**:  https://arxiv.org/abs/2401.04088")
            
            gr.Markdown("# 📖 𝓬𝓼 𝓪𝓻𝓧𝓲𝓿: 𝟮𝟰𝟬𝟭.𝟬𝟰𝟬𝟴𝟴 💫")
            gr.Markdown("**arxiv**:  https://arxiv.org/pdf/2401.04088")


            story_starters = gr.Dataframe(
                value=STORY_STARTERS,
                headers=["Category", "Starter"],
                interactive=False
            )
            
            gr.Markdown("### 🎬 Gallery")
            gallery = gr.Dataframe(
                value=load_gallery(),
                headers=["Timestamp", "Preview", "Story Path", "Audio Path"],
                interactive=False
            )

            gr.Markdown("""  
            
## Understanding Mixtral's Expert Layers Per Claude 3.5 Sonnet

### What is Mixtral?
Imagine having 8 different expert tutors for each topic you're learning. Instead of consulting all of them every time, you pick the 2 best experts for each specific question. That's basically how Mixtral works! 

### Key Concepts Broken Down

#### 1. The Expert System
- Total Experts: 8 specialists per layer
- Active Experts: 2 chosen per task
- Total Layers: 32 levels of processing
- Selection Method: Uses a "router" to pick the best experts

#### 2. How It Works (Step by Step)

1. **Input Stage**
   - Text enters the system
   - Gets broken down into tokens (like words or parts of words)

2. **Router Selection**
   - A "smart traffic director" (router) looks at each piece of input
   - Chooses 2 out of 8 experts to handle it
   - Makes selections based on what each expert is best at

3. **Expert Processing**
   - Chosen experts analyze the input
   - Each expert contributes their specialty knowledge
   - Results are combined using weighted averages

4. **Pattern Analysis**
   Layer 0 (Beginning): 
   - Acts almost randomly (~14% repeat choices)
   - Like students still finding their preferred study partners

   Layer 15 (Middle):
   - Shows stronger patterns (~27% repeat choices)
   - Experts develop clear specialties

   Layer 31 (End):
   - Maintains expertise patterns (~22% repeat choices)
   - Balanced between specialization and flexibility

#### 3. Expert Choice Patterns (From Real Data)

| Location in Network | Expert Reuse Rate | What It Means |
|-------------------|-------------------|---------------|
| Early Layers | 12.5-14% | Nearly random selection |
| Middle Layers | 25-28% | Strong specialization |
| Final Layers | 20-25% | Balanced expertise |

#### 4. Why This Matters

- **Efficiency**: Only uses 2/8 experts at a time = faster processing
- **Specialization**: Experts become really good at specific tasks
- **Flexibility**: Different expert combinations = handles varied tasks well
- **Resource Smart**: Like having 8 teachers but only paying for 2 at a time!

## Real-World Example
Think of it like a hospital:
- You have 8 specialists available
- For each patient, you consult the 2 most relevant doctors
- Your choices get better as you learn which doctors handle which cases best
- Sometimes you need the same doctors for related cases (explaining the repeat patterns)

*Note: The actual repeat rates vary by data type (math, code, text, etc.) showing how different experts specialize in different areas!*


---


## Explanation of the Eight-Layer Expert Weighting in Mixtral Per ChatGPT 4o Omni Model


Mixtral 8x7B is a sparse Mixture of Experts (MoE) model where each layer can choose from eight "experts" or sets of specialized parameters. For each token, Mixtral’s routing system selects two experts from these eight at every layer, reducing computational cost by using only part of its parameters per token. This selective process means that different experts are activated for different types of input, allowing for efficient and context-sensitive processing.

Key concepts:

Experts: Each expert is a set of parameters trained on specific types of data, like code or medical abstracts. These experts specialize, enhancing their ability to process particular kinds of input.
Layers and Repetition: Layers 15 and 31 show higher "expert assignment repetitions" (temporal locality), meaning certain experts are frequently reselected across consecutive tokens, likely due to stable contextual relevance at these stages.
First and Second Choices: The system tracks the first and second choice for each expert across layers, where layers 15 and 31 have distinct patterns of assignment compared to random, as evidenced by percentages surpassing expected random choice (12.5%).

# Mixtral 8x7B Layer and Expert Weighting Overview

The Mixtral 8x7B model is built on an efficient and scalable design where only a subset of its eight expert networks, or "experts," are activated at each layer. This selective routing allows Mixtral to use fewer parameters per token while maintaining high performance, making it faster and more efficient.

| **Layer** | **First Choice Repetition (%)** | **First or Second Choice Repetition (%)** | **Explanation**                                                                                       |
|-----------|----------------------------------|-------------------------------------------|-------------------------------------------------------------------------------------------------------|
| Layer 0   | ~14.0 - 15.0%                    | ~46 - 50%                                 | Initial layers have near-random expert selection, establishing basic token processing.               |
| Layer 15  | 23.6 - 28.4%                     | 61.6 - 67.0%                              | This layer shows a peak in expert repetition, enhancing stable context for language understanding.   |
| Layer 31  | 19.7 - 26.3%                     | 44.5 - 53.6%                              | Final layers focus on consolidating information, applying specific experts based on accumulated context. |

### Method Steps
1. **Token Processing**: Each token in a sequence is routed through a "layer," which selects two experts from eight possibilities to process the token.
2. **Expert Selection**: At each layer, Mixtral uses a gating network to pick two experts. The choice is influenced by the token's properties and the layer's context.
3. **Repetition and Efficiency**: Layers 15 and 31 have high expert repetition rates. This repetition is not random and suggests the model’s experts are increasingly stable, applying consistent expertise as it processes deeper into a sequence.
4. **Interpretation**: High expert repetition means Mixtral is refining its processing. It returns to the same experts for similar content, building a coherent response across consecutive tokens.
5. **Output Generation**: The experts chosen at each layer contribute to the final response, weighted by how frequently they’re selected across the layers.

By selecting only a subset of experts per token and layer, Mixtral achieves both computational efficiency and high contextual relevance across tokens in the input sequence.

### Summary of Table Data
This Markdown table concisely summarizes Mixtral’s expert assignment behavior across layers and explains the significance of the repetition patterns observed. 

This structure offers clarity on how Mixtral leverages expert assignments, especially in layers 15 and 31, for efficient processing while retaining contextual understanding!

""")
    
    # Event handlers
    def update_prompt(evt: gr.SelectData):
        return STORY_STARTERS[evt.index[0]][1]
    
    story_starters.select(update_prompt, None, prompt_input)
    
    generate_btn.click(
        fn=process_story_and_audio,
        inputs=[prompt_input, model_choice],
        outputs=[story_output, audio_output, gallery]
    )
    
    gallery.select(
        fn=play_gallery_audio,
        inputs=[gallery],
        outputs=[audio_output]
    )

if __name__ == "__main__":
    demo.launch()