app.py

import os
import gradio as gr
import openai
import PyPDF2
import numpy as np
import math

MODEL_STATUS = {
    'tiktoken': False,
    'transformers': False,
    'torch': False,
    'model_loaded': False,
    'error_messages': []
}

try:
    import tiktoken
    gpt_tokenizer = tiktoken.get_encoding("gpt2")
    MODEL_STATUS['tiktoken'] = True
except Exception as e:
    MODEL_STATUS['error_messages'].append(f"tiktoken error: {str(e)}")
    gpt_tokenizer = None

# WEEK 3 
# try:
#     from transformers import AutoTokenizer, AutoModel
#     import torch
#     MODEL_STATUS['transformers'] = True
#     MODEL_STATUS['torch'] = True
#     
#     print("Loading model...")
#     tokenizer = AutoTokenizer.from_pretrained("prajjwal1/bert-tiny")
#     model = AutoModel.from_pretrained("prajjwal1/bert-tiny")
#     MODEL_STATUS['model_loaded'] = True
#     print("model loaded successfully!")
#     
# except Exception as e:
#     MODEL_STATUS['error_messages'].append(f"Model loading error: {str(e)}")
#     tokenizer = None
#     model = None

tokenizer = None
model = None

# OpenAI setup
OPENAI_API_KEY = os.getenv("openAI_TOKEN")
if OPENAI_API_KEY:
    openai.api_key = OPENAI_API_KEY
else:
    MODEL_STATUS['error_messages'].append("OpenAI API key not found")

import shutil
import os

cache_dir = os.path.expanduser("~/.cache/huggingface")
if os.path.exists(cache_dir):
    try:
        total_size = sum(
            os.path.getsize(os.path.join(dirpath, filename))
            for dirpath, dirnames, filenames in os.walk(cache_dir)
            for filename in filenames
        ) / (1024**3)  
        
        if total_size > 40: 
            shutil.rmtree(cache_dir)
            print(f"Cleared {total_size:.2f}GB cache")
    except Exception as e:
        print(f"Cache cleanup error: {e}")

from model_functions import *

def tokenize_text(text):
    if not text.strip():
        return [], 0, "Enter some text to see tokenization"
    
    if gpt_tokenizer:
        try:
            tokens = gpt_tokenizer.encode(text)
            token_strings = []
            for token in tokens:
                try:
                    decoded = gpt_tokenizer.decode([token])
                    token_strings.append(decoded)
                except UnicodeDecodeError:
                    token_strings.append(f"<token_{token}>")
            return token_strings, len(tokens), f"Text tokenized successfully → {len(tokens)} tokens"
        except Exception as e:
            return [], 0, f"Tokenization error: {str(e)}"
    else:
        # Fallback: simple whitespace tokenization
        tokens = text.split()
        return tokens, len(tokens), f"Using fallback tokenization → {len(tokens)} tokens (tiktoken unavailable)"

def get_next_token_predictions(text):
    """Get next token predictions using OpenAI API"""
    if not text.strip():
        return "Enter some text to see predictions"
    
    if not OPENAI_API_KEY:
        return "OpenAI API key not available - cannot generate predictions"
    
    try:
        client = openai.OpenAI(api_key=OPENAI_API_KEY)
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "Complete the following text with the next most likely word. Provide exactly 3 options with their approximate probabilities."},
                {"role": "user", "content": f"Text: '{text}'\n\nNext word options:"}
            ],
            temperature=0.1,
            max_tokens=50
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Error getting predictions: {str(e)}"

def merge_subword_tokens(tokens, attention_matrix):
    """Merge subword tokens back into words for cleaner viz"""
    merged_tokens = []
    merged_attention = []
    current_word = ""
    current_indices = []
    
    for i, token in enumerate(tokens):
        if token.startswith('##'):
            current_word += token[2:]
            current_indices.append(i)
        else:
            if current_word:
                merged_tokens.append(current_word)
                merged_attention.append(current_indices)
            current_word = token
            current_indices = [i]
    
    if current_word:
        merged_tokens.append(current_word)
        merged_attention.append(current_indices)
    
    # Merge attention weights by averaging
    merged_matrix = np.zeros((len(merged_tokens), len(merged_tokens)))
    for i, i_indices in enumerate(merged_attention):
        for j, j_indices in enumerate(merged_attention):
            # Average attention between word groups
            weights = []
            for ii in i_indices:
                for jj in j_indices:
                    if ii < attention_matrix.shape[0] and jj < attention_matrix.shape[1]:
                        weights.append(attention_matrix[ii, jj])
            if weights:
                merged_matrix[i, j] = np.mean(weights)
    
    return merged_tokens, merged_matrix

def create_attention_network_svg(text):
    if not text.strip():
        return "Enter text to see attention network"
    
    if not MODEL_STATUS['model_loaded']:
        return f"Attention model not available. Errors: {MODEL_STATUS['error_messages']}"

    try:
        # Tokenize input
        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=64)
        tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])

        with torch.no_grad():
            outputs = model(**inputs, output_attentions=True)

        # Remove special tokens
        clean_tokens = []
        clean_indices = []
        for i, token in enumerate(tokens):
            if token not in ['[CLS]', '[SEP]', '[PAD]']:
                clean_tokens.append(token)
                clean_indices.append(i)

        if len(clean_indices) < 2:
            return "Need at least 2 valid tokens for attention visualisation."

        # SEARCH for best head: max variance
        best_attention = None
        best_name = ""
        best_tokens = []
        best_variance = -1

        debug_info = f"Total Layers: {len(outputs.attentions)}\n"

        for layer_idx, layer_att in enumerate(outputs.attentions):
            num_heads = layer_att.shape[1]
            for head_idx in range(num_heads):
                attn_matrix = layer_att[0, head_idx].numpy()
                trimmed_attention = attn_matrix[np.ix_(clean_indices, clean_indices)]
                variance = np.var(trimmed_attention)

                debug_info += f"Layer {layer_idx}, Head {head_idx} — Variance: {variance:.5f}\n"

                if variance > best_variance:
                    best_attention = trimmed_attention
                    best_name = f"Layer {layer_idx}, Head {head_idx}"
                    best_tokens = clean_tokens
                    best_variance = variance

        if best_attention is None:
            return "Could not extract valid attention."

        # Merge subwords
        merged_tokens, merged_attention = merge_subword_tokens(best_tokens, best_attention)
        n_tokens = len(merged_tokens)

        if n_tokens < 2:
            return "Too few tokens after merging for attention graph."

        # SVG dimensions
        width, height = 1000, 500
        margin = 50

        # Linear positions
        positions = []
        for i in range(n_tokens):
            x = margin + (width - 2*margin) * i / (n_tokens - 1)
            y = height // 2
            positions.append((x, y))

        # Start SVG
        svg = f'<svg width="{width}" height="{height}" xmlns="http://www.w3.org/2000/svg">'
        svg += '<style>.token-text { font-family: Arial; font-size: 14px; text-anchor: middle; font-weight: bold; }'
        svg += '.debug-text { font-family: monospace; font-size: 10px; fill: #666; }</style>'

        # Choose top-N attention connections
        num_top_connections = 20
        pairs = []
        for i in range(n_tokens):
            for j in range(n_tokens):
                if i != j:
                    pairs.append((merged_attention[i, j], i, j))
        pairs.sort(reverse=True)
        top_pairs = pairs[:num_top_connections]

        # Draw attention arcs
        for weight, i, j in top_pairs:
            x1, y1 = positions[i]
            x2, y2 = positions[j]
            mid_x = (x1 + x2) / 2
            curve_y = y1 - 80 if (i + j) % 2 == 0 else y1 + 80

            # Color coding
            if weight > 0.08:
                color = "#d32f2f"  # red
                opacity = "0.8"
            elif weight > 0.04:
                color = "#ff9800"  # orange
                opacity = "0.6"
            else:
                color = "#2196f3"  # blue
                opacity = "0.4"

            thickness = max(2, weight * 10)

            svg += f'<path d="M {x1},{y1} Q {mid_x},{curve_y} {x2},{y2}" '
            svg += f'stroke="{color}" stroke-width="{thickness}" fill="none" opacity="{opacity}"/>'

        # Draw nodes
        for i, (token, (x, y)) in enumerate(zip(merged_tokens, positions)):
            svg += f'<circle cx="{x}" cy="{y}" r="25" fill="white" stroke="black" stroke-width="2"/>'
            svg += f'<text x="{x}" y="{y+5}" class="token-text">{token[:10]}</text>'

        # Legend and info
        svg += f'<text x="20" y="{height - 130}" style="font-family: Arial; font-size: 16px; font-weight: bold;">'
        svg += f'Attention Network - {best_name}</text>'
        svg += f'<text x="20" y="{height - 110}" style="font-family: Arial; font-size: 12px;">'
        svg += f'Red: Strong | Orange: Medium | Blue: Weak | Showing top {num_top_connections} connections</text>'

        # Debug info (limited lines)
        for i, line in enumerate(debug_info.split('\n')[:8]):
            svg += f'<text x="20" y="{height - 90 + 12*i}" class="debug-text">{line}</text>'

        svg += '</svg>'

        return svg

    except Exception as e:
        return f"Error generating attention network: {str(e)}"


with gr.Blocks() as demo:
    gr.Markdown("# Language Models & Methods Lab Interface")
    
   
    with gr.Tabs() as tabs:
        
        # Week 3 Tab
        with gr.Tab("Week 3: Text Processing"):
            gr.Markdown("# How Language Models Process Text")
            gr.Markdown("Explore tokenization, context windows, and attention mechanisms")
            
            with gr.Tabs() as week3_tabs:
                with gr.Tab("Tokenization Explorer"):
                    gr.Markdown("### See how text gets broken into tokens")
                    
                    with gr.Row():
                        token_input = gr.Textbox(
                            label="Enter your text",
                            placeholder="Type any text to see how it gets tokenized...",
                            lines=3,
                            value="The quick brown fox jumps over the lazy dog."
                        )
                    
                    with gr.Row():
                        tokenize_btn = gr.Button("Tokenize Text")
                    
                    with gr.Row():
                        token_display = gr.Textbox(label="Tokens", lines=3, interactive=False)
                        token_count = gr.Number(label="Token Count", interactive=False)
                    
                    with gr.Row():
                        token_info = gr.Textbox(label="Tokenization Info", lines=2, interactive=False)
                
                with gr.Tab("Context & Predictions"):
                    gr.Markdown("### Next-word predictions and context understanding")
                    
                    with gr.Row():
                        context_input = gr.Textbox(
                            label="Enter incomplete text",
                            placeholder="I went to the bank to",
                            lines=2,
                            value="I went to the bank to"
                        )
                    
                    with gr.Row():
                        predict_btn = gr.Button("Get Next Word Predictions")
                    
                    with gr.Row():
                        predictions_output = gr.Textbox(label="Most Likely Next Words", lines=5, interactive=False)
                    
                    with gr.Row():
                        context_window_info = gr.Textbox(
                            label="Context Window Status",
                            value="Click 'Get Predictions' to see token usage",
                            interactive=False
                        )
                
                with gr.Tab("Attention Network"):
                    gr.Markdown("### Network visualisation of attention patterns")
                    gr.Markdown("See how words connect to each other through attention mechanisms")
                    
                    with gr.Row():
                        attention_input = gr.Textbox(
                            label="Enter a sentence (shorter sentences work better)",
                            placeholder="The bank was closed.",
                            lines=2,
                            value="The bank was closed."
                        )
                    
                    with gr.Row():
                        analyze_attention_btn = gr.Button("Generate Attention Network")
                    
                    with gr.Row():
                        attention_network = gr.HTML(label="Attention Network Visualisation")
            
            # Week 3 Event Handlers
            def update_tokenization(text):
                tokens, count, info = tokenize_text(text)
                token_str = " | ".join(tokens) if tokens else ""
                return token_str, count, info
            
            def update_predictions_with_context(text):
                if not text.strip():
                    return "Enter text to get predictions", "No text to analyze"
                
                # Get token count for context window
                _, token_count, _ = tokenize_text(text)
                context_status = f"Current text: {token_count} tokens / 4096 (GPT-3.5 limit) = {token_count/4096*100:.1f}% used"
                
                # Get predictions
                predictions = get_next_token_predictions(text)
                
                return predictions, context_status
            
            def generate_network_visualization(text):
                return create_attention_network_svg(text)
            
            # Connect event handlers
            tokenize_btn.click(
                update_tokenization,
                inputs=[token_input],
                outputs=[token_display, token_count, token_info]
            )
            
            # Auto-update tokenization as user types
            token_input.change(
                update_tokenization,
                inputs=[token_input],
                outputs=[token_display, token_count, token_info]
            )
            
            predict_btn.click(
                update_predictions_with_context,
                inputs=[context_input],
                outputs=[predictions_output, context_window_info]
            )
            
            analyze_attention_btn.click(
                generate_network_visualization,
                inputs=[attention_input],
                outputs=[attention_network]
            )
        
        # OTHER WEEKS

        with gr.Tab("Week 4: Controlling Model Behaviour"):
            gr.Markdown("# Controlling Model Behaviour Through Prompting")
            gr.Markdown("Explore how different prompting techniques and parameters affect model outputs")
            
            with gr.Tabs() as week4_tabs:
                
                with gr.Tab("Temperature Effects"):
                    gr.Markdown("### Compare how temperature affects creativity and consistency")
                    
                    with gr.Row():
                        temp_input = gr.Textbox(
                            label="Enter your prompt",
                            placeholder="Type your question or prompt here...",
                            lines=3,
                            value="Write a creative opening sentence for a story about a robot looking for a friend."
                        )
                    
                    with gr.Row():
                        temp_slider1 = gr.Slider(
                            minimum=0.1, 
                            maximum=0.4, 
                            value=0.2, 
                            step=0.1, 
                            label="Low Temperature (More Focused & Consistent)"
                        )
                        temp_slider2 = gr.Slider(
                            minimum=0.7, 
                            maximum=1.0, 
                            value=0.9, 
                            step=0.1, 
                            label="High Temperature (More Creative & Varied)"
                        )

                    with gr.Row():
                        generate_temp = gr.Button("Generate Both Responses")
                    
                    with gr.Row():
                        focused_output = gr.Textbox(
                            label="Focused Output (Low Temperature)", 
                            lines=5
                        )
                        creative_output = gr.Textbox(
                            label="Creative Output (High Temperature)", 
                            lines=5
                        )
                
                with gr.Tab("System Prompts"):
                    gr.Markdown("### See how system prompts shape model behaviour")
                    
                    with gr.Row():
                        system_input = gr.Textbox(
                            label="Enter your prompt",
                            placeholder="Type your question or prompt here...",
                            lines=3,
                            value="Explain what a database index is."
                        )
                    
                    with gr.Row():
                        system_prompt_dropdown = gr.Dropdown(
                            choices=[
                                "You are a helpful assistant providing accurate, concise answers.",
                                "You are a data scientist explaining technical concepts with precision and examples.",
                                "You are a creative storyteller who uses vivid metaphors and analogies.",
                                "You are a critical reviewer who evaluates information carefully and points out limitations.",
                                "You are a friendly teacher explaining concepts to someone learning for the first time."
                            ],
                            label="Choose System Prompt",
                            value="You are a helpful assistant providing accurate, concise answers."
                        )
                    
                    with gr.Row():
                        generate_system = gr.Button("Generate Response")
                    
                    with gr.Row():
                        system_output = gr.Textbox(label="Output", lines=6)
                
                with gr.Tab("Prompting Techniques"):
                    gr.Markdown("""
                    ### Compare Zero-Shot, Few-Shot, and Chain-of-Thought
                    - **Zero-shot:** Direct question without examples
                    - **Few-shot:** You should provide similar examples to guide the response
                    - **Chain-of-thought:** Asks model to break down reasoning step-by-step
                    """)
                    
                    with gr.Row():
                        shot_input = gr.Textbox(
                            label="Enter your task",
                            placeholder="Enter a task that requires reasoning...",
                            lines=3,
                            value="Classify the sentiment: 'The product works okay but customer service was terrible.'"
                        )
                    
                    with gr.Row():
                        approach_type = gr.Radio(
                            ["zero-shot", "few-shot", "chain-of-thought"],
                            label="Select Prompting Technique",
                            value="zero-shot"
                        )
                    
                    with gr.Row():
                        generate_shot = gr.Button("Generate Response")
                    
                    with gr.Row():
                        shot_output = gr.Textbox(label="Output", lines=8)
                
                with gr.Tab("Combining Techniques"):
                    gr.Markdown("### Experiment with combining multiple techniques")
                    
                    with gr.Row():
                        combo_input = gr.Textbox(
                            label="Enter your task",
                            placeholder="Enter a complex task...",
                            lines=3,
                            value="Analyse this review and suggest improvements: 'App crashes sometimes but has good features.'"
                        )
                    
                    with gr.Row():
                        combo_system = gr.Dropdown(
                            choices=[
                                "None (default)",
                                "You are a product analyst providing structured feedback.",
                                "You are a UX researcher focused on user experience.",
                            ],
                            label="System Prompt (optional)",
                            value="None (default)"
                        )
                    
                    with gr.Row():
                        combo_examples = gr.Checkbox(
                            label="Include few-shot examples",
                            value=False
                        )
                        combo_cot = gr.Checkbox(
                            label="Use chain-of-thought reasoning",
                            value=False
                        )
                    
                    with gr.Row():
                        combo_temp = gr.Slider(
                            minimum=0.1,
                            maximum=1.0,
                            value=0.5,
                            step=0.1,
                            label="Temperature"
                        )
                    
                    with gr.Row():
                        generate_combo = gr.Button("Generate Response")
                    
                    with gr.Row():
                        combo_output = gr.Textbox(label="Output", lines=8)
                        combo_info = gr.Textbox(label="Techniques Applied", lines=4)

            generate_temp.click(
                lambda x, t1, t2: [
                    generate_with_temperature(x, t1), 
                    generate_with_temperature(x, t2)
                ],
                inputs=[temp_input, temp_slider1, temp_slider2],
                outputs=[focused_output, creative_output]
            )

            generate_system.click(
                generate_with_system_prompt,
                inputs=[system_input, system_prompt_dropdown],
                outputs=system_output
            )

            generate_shot.click(
                generate_with_examples,
                inputs=[shot_input, approach_type],
                outputs=shot_output
            )

            generate_combo.click(
                generate_combined_techniques,
                inputs=[combo_input, combo_system, combo_examples, combo_cot, combo_temp],
                outputs=[combo_output, combo_info]
            )
        
        with gr.Tab("Week 5: Advanced Prompting"):
            gr.Markdown("# Advanced Prompt Engineering Techniques")
            gr.Markdown("Explore sophisticated prompting strategies and visualise reasoning patterns")
            
            with gr.Tabs() as week5_tabs:
                
                with gr.Tab("Tree of Thought Explorer"):
                    gr.Markdown("""
                    ### Visualise Multi-Path Reasoning
                    The model will break down your problem into multiple approaches, evaluate each one, and select the best path.
                    """)
                    
                    with gr.Row():
                        tot_input = gr.Textbox(
                            label="Enter a problem to solve",
                            placeholder="e.g., How can we improve user engagement on a mobile app?",
                            lines=3,
                            value="How should a startup decide between building a mobile app or a web application first?"
                        )
                    
                    with gr.Row():
                        generate_tot = gr.Button("Generate Tree of Thought", variant="primary")
                    
                    with gr.Row():
                        tot_output = gr.Textbox(
                            label="Reasoning Process",
                            lines=12
                        )
                    
                    with gr.Row():
                        tot_visualization = gr.HTML(
                            label="Tree Visualisation"
                        )
                
                with gr.Tab("Self-Consistency Testing"):
                    gr.Markdown("""
                    ### Test Response Consistency
                    Run the same prompt multiple times to identify consistent patterns and areas of uncertainty.
                    """)
                    
                    with gr.Row():
                        consistency_input = gr.Textbox(
                            label="Enter your prompt",
                            placeholder="Ask a question that requires reasoning...",
                            lines=3,
                            value="What are the three most important factors in choosing a database system?"
                        )
                    
                    with gr.Row():
                        num_runs = gr.Slider(
                            minimum=3,
                            maximum=5,
                            value=3,
                            step=1,
                            label="Number of generations"
                        )
                        consistency_temp = gr.Slider(
                            minimum=0.3,
                            maximum=0.9,
                            value=0.7,
                            step=0.1,
                            label="Temperature"
                        )
                    
                    with gr.Row():
                        generate_consistency = gr.Button("Generate Multiple Responses", variant="primary")
                    
                    with gr.Row():
                        consistency_analysis = gr.Textbox(
                            label="Analysis Guide",
                            lines=4
                        )
                    
                    with gr.Row():
                        consistency_output1 = gr.Textbox(label="Response 1", lines=5)
                        consistency_output2 = gr.Textbox(label="Response 2", lines=5)
                    
                    with gr.Row():
                        consistency_output3 = gr.Textbox(label="Response 3", lines=5)
                        consistency_output4 = gr.Textbox(label="Response 4 (if selected)", lines=5, visible=True)
                    
                    with gr.Row():
                        consistency_output5 = gr.Textbox(label="Response 5 (if selected)", lines=5, visible=True)
                
                with gr.Tab("Prompt Structure Comparison"):
                    gr.Markdown("""
                    ### Compare Structural Strategies
                    Test how different prompt structures affect model attention and output quality.
                    """)
                    
                    with gr.Row():
                        structure_input = gr.Textbox(
                            label="Enter your task",
                            placeholder="Enter a task or question...",
                            lines=3,
                            value=""
                        )
                    
                    with gr.Row():
                        gr.Markdown("### Select ONE structure to test:")
                    
                    with gr.Row():
                        structure_radio = gr.Radio(
                            choices=[
                                "Baseline (no special structure)",
                                "Front-loading (critical instruction first)",
                                "Delimiter strategy (section separation)",
                                "Sandwich technique (instruction at start and end)"
                            ],
                            label="Prompt Structure",
                            value="Baseline (no special structure)"
                        )
                    
                    with gr.Row():
                        generate_structure = gr.Button("Generate Response", variant="primary")
                    
                    with gr.Row():
                        structure_output = gr.Textbox(
                            label="Response",
                            lines=8
                        )
                        structure_info = gr.Textbox(
                            label="Structure Information",
                            lines=8
                        )
            
            # Week 5 Event Handlers
            def handle_tot(task):
                text_output, svg_output = generate_tot_response(task)
                return text_output, svg_output
            
            def handle_consistency(prompt, runs, temp):
                responses, analysis = generate_self_consistency(prompt, int(runs), temp)
                while len(responses) < 5:
                    responses.append("")
                return analysis, responses[0], responses[1], responses[2], responses[3], responses[4]
            
            def handle_structure(task, structure_choice):
                use_frontload = "Front-loading" in structure_choice
                use_delimiters = "Delimiter" in structure_choice
                use_sandwich = "Sandwich" in structure_choice
                
                output, info = compare_prompt_structures(task, use_frontload, use_delimiters, use_sandwich)
                return output, info
            
            generate_tot.click(
                handle_tot,
                inputs=[tot_input],
                outputs=[tot_output, tot_visualization]
            )
            
            generate_consistency.click(
                handle_consistency,
                inputs=[consistency_input, num_runs, consistency_temp],
                outputs=[consistency_analysis, consistency_output1, consistency_output2, 
                        consistency_output3, consistency_output4, consistency_output5]
            )
            
            generate_structure.click(
                handle_structure,
                inputs=[structure_input, structure_radio],
                outputs=[structure_output, structure_info]
            )


        
        
        
        with gr.Tab("Assignment 1"):
            gr.Markdown("# Assignment 1: Prompting Strategy Evaluation")
            gr.Markdown("""
            Test different prompting strategies for your chosen NLP task.
            Remember: You need 3 documents, with 2 different strategies tested per document (6 total experiments).
            """)
            
            with gr.Row():
                assignment_task = gr.Dropdown(
                    choices=["Sentiment Analysis", "Summarisation"],
                    label="Select NLP Task",
                    value="Sentiment Analysis"
                )
            
            with gr.Row():
                with gr.Column():
                    assignment_text = gr.Textbox(
                        label="Enter Text",
                        placeholder="Paste your document text here...",
                        lines=6
                    )
                with gr.Column():
                    assignment_file = gr.File(
                        label="OR Upload a File (TXT or PDF)",
                        file_types=[".txt", ".pdf"],
                        type="binary"
                    )
            
            gr.Markdown("### Select Your Prompting Strategy")
            
            with gr.Row():
                strategy_type = gr.Radio(
                    choices=[
                        "Direct (no special technique)",
                        "Chain-of-thought (step-by-step reasoning)",
                        "Role-based (uses system prompt)",
                        "Combined (role + chain-of-thought)"
                    ],
                    label="Prompting Strategy",
                    value="Direct (no special technique)",
                    info="Choose how the model should approach the task"
                )
            
            with gr.Row():
                system_role = gr.Dropdown(
                    choices=[
                        "None",
                        "Technical analyst",
                        "Creative assistant"
                    ],
                    label="System Role (for role-based strategies)",
                    value="None",
                    info="Only applies if you selected a role-based strategy"
                )
            
            with gr.Row():
                assignment_temp = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.5,
                    step=0.1,
                    label="Temperature (0.1 = focused, 1.0 = creative)"
                )
            
            with gr.Row():
                generate_assignment = gr.Button("Generate Response", variant="primary")
            
            with gr.Row():
                assignment_output = gr.Textbox(
                    label="Model Output",
                    lines=12
                )
            
            with gr.Row():
                assignment_info = gr.Textbox(
                    label="Strategy Applied",
                    lines=3,
                    info="Documents which settings were used for this experiment"
                )

            generate_assignment.click(
                handle_assignment_experiment,
                inputs=[assignment_text, assignment_file, assignment_task, strategy_type, system_role, assignment_temp],
                outputs=[assignment_output, assignment_info]
            )
        # with gr.Tab("Week 8: Error Detection"):
        #     # Week 8 content here
        #     pass

        with gr.Tab("Week 9: Evaluation & Quality Assessment"):
            gr.Markdown("# Evaluation & Quality Assessment")
            gr.Markdown("Practice evaluating LLM outputs using techniques from today's lecture")
            
            with gr.Tabs() as week9_tabs:
                
                with gr.Tab("Human Evaluation"):
                    gr.Markdown("""
                    ### Generate Multiple Versions for Comparison
                    Create three versions of a response with different temperature settings. 
                    Then rate each output to practice human evaluation.
                    """)
                    
                    with gr.Row():
                        eval_prompt = gr.Textbox(
                            label="Enter your prompt",
                            placeholder="e.g., Summarise the main benefits of cloud computing for small businesses",
                            lines=3,
                            value="Write three different, creative metaphors to explain the concept of a neural network to a child."
                        )
                    
                    with gr.Row():
                        with gr.Column():
                            gr.Markdown("**Temperature Settings:**")
                            eval_temp1 = gr.Slider(
                                minimum=0.1,
                                maximum=1.0,
                                value=0.3,
                                step=0.1,
                                label="Version 1 Temperature (Focused)"
                            )
                            eval_temp2 = gr.Slider(
                                minimum=0.1,
                                maximum=1.0,
                                value=0.7,
                                step=0.1,
                                label="Version 2 Temperature (Balanced)"
                            )
                            eval_temp3 = gr.Slider(
                                minimum=0.1,
                                maximum=1.0,
                                value=1.0,
                                step=0.1,
                                label="Version 3 Temperature (Creative)"
                            )
                    
                    with gr.Row():
                        generate_eval_btn = gr.Button("Generate 3 Versions", variant="primary")
                    
                    gr.Markdown("### Compare and Rate the Outputs")
                    
                    with gr.Row():
                        with gr.Column():
                            gr.Markdown("**Version 1** (Temp: 0.3)")
                            eval_output1 = gr.Textbox(
                                label="Output 1",
                                lines=6,
                                # interactive=False
                            )
                            gr.Markdown("**Rate this output (1=Poor, 5=Excellent):**")
                            with gr.Row():
                                rate1_accuracy = gr.Slider(1, 5, value=3, step=1, label="Accuracy")
                                rate1_coherence = gr.Slider(1, 5, value=3, step=1, label="Coherence")
                            with gr.Row():
                                rate1_completeness = gr.Slider(1, 5, value=3, step=1, label="Completeness")
                                rate1_relevance = gr.Slider(1, 5, value=3, step=1, label="Relevance")
                        
                        with gr.Column():
                            gr.Markdown("**Version 2** (Temp: 0.7)")
                            eval_output2 = gr.Textbox(
                                label="Output 2",
                                lines=6,
                                # interactive=False
                            )
                            gr.Markdown("**Rate this output (1=Poor, 5=Excellent):**")
                            with gr.Row():
                                rate2_accuracy = gr.Slider(1, 5, value=3, step=1, label="Accuracy")
                                rate2_coherence = gr.Slider(1, 5, value=3, step=1, label="Coherence")
                            with gr.Row():
                                rate2_completeness = gr.Slider(1, 5, value=3, step=1, label="Completeness")
                                rate2_relevance = gr.Slider(1, 5, value=3, step=1, label="Relevance")
                        
                        with gr.Column():
                            gr.Markdown("**Version 3** (Temp: 1.0)")
                            eval_output3 = gr.Textbox(
                                label="Output 3",
                                lines=6,
                                # interactive=False
                            )
                            gr.Markdown("**Rate this output (1=Poor, 5=Excellent):**")
                            with gr.Row():
                                rate3_accuracy = gr.Slider(1, 5, value=3, step=1, label="Accuracy")
                                rate3_coherence = gr.Slider(1, 5, value=3, step=1, label="Coherence")
                            with gr.Row():
                                rate3_completeness = gr.Slider(1, 5, value=3, step=1, label="Completeness")
                                rate3_relevance = gr.Slider(1, 5, value=3, step=1, label="Relevance")
                    
                    with gr.Row():
                        calculate_ratings_btn = gr.Button("Calculate Average Ratings")
                    
                    with gr.Row():
                        ratings_summary = gr.Textbox(
                            label="Ratings Summary",
                            lines=6,
                            # interactive=False
                        )
                with gr.Tab("Automatic Evaluation"):
                    gr.Markdown("""
                    ### Generate a Response and Compare to Your Reference Answer
                    
                    This demonstrates how automatic metrics like BLEU and word overlap work in practice.
                    You'll provide a "reference answer" (what a good response should say), then see how
                    the model's response compares using automatic metrics.
                    """)
                    
                    with gr.Row():
                        metric_prompt = gr.Textbox(
                            label="Enter your prompt (question or task)",
                            placeholder="e.g., What are the main benefits of using a relational database?",
                            lines=3,
                            value="What are the three main principles of user-centered design?"
                        )
                    
                    with gr.Row():
                        metric_reference = gr.Textbox(
                            label="Enter your reference answer (what a good answer should include)",
                            placeholder="Write what you consider a good/correct answer to your prompt...",
                            lines=5,
                            value="The three main principles of user-centered design are: 1) Focus on users and their needs throughout the design process, 2) Involve users early and often through testing and feedback, and 3) Iterate designs based on user feedback to continuously improve the experience."
                        )
                    
                    with gr.Row():
                        metric_temp = gr.Slider(
                            minimum=0.1,
                            maximum=1.0,
                            value=0.7,
                            step=0.1,
                            label="Temperature"
                        )
                    
                    with gr.Row():
                        generate_metric_btn = gr.Button("Generate Model Response & Calculate Metrics", variant="primary")
                    
                    gr.Markdown("### Model Response")
                    
                    with gr.Row():
                        metric_generated = gr.Textbox(
                            label="Generated Answer (model's response)",
                            lines=6,
                            # interactive=False
                        )
                    
                    gr.Markdown("### Evaluation Metrics")
                    
                    with gr.Row():
                        with gr.Column():
                            metric_overlap_display = gr.Textbox(
                                label="Word Overlap",
                                lines=1,
                                # interactive=False
                            )
                        with gr.Column():
                            gr.Markdown("**Quick Summary:** This shows the % of reference words that appear in the generated response")
                    
                    with gr.Row():
                        metric_report = gr.Textbox(
                            label="Detailed Metrics Report",
                            lines=18,
                            # interactive=False
                        )
                    
                    gr.Markdown("""
                    ### Understanding the Metrics
                    
                    **Word Overlap:** What % of words from your reference appear in the generated response?
                    - Shows which words matched, which were missing, which were added
                    - High overlap = similar vocabulary used
                    
                    **BLEU Score:** Modified word overlap that penalises very short responses
                    - Used commonly for translation and summarisation
                    - Ranges roughly 0-100 (higher = more overlap)
                    
                    **Important Limitations:**
                    - These metrics only measure word overlap, NOT meaning or quality
                    - A response with low overlap might still be correct (using synonyms)
                    - A response with high overlap might still be wrong (same words, wrong meaning)
                    - Always use human judgment alongside automatic metrics!
                    """)
                
            
            def update_consistency_visibility(num_runs):
                """Show/hide output boxes based on number of runs"""
                num_runs = int(num_runs)
                return (
                    gr.update(visible=True),  # output1 always visible
                    gr.update(visible=True),  # output2 always visible
                    gr.update(visible=True),  # output3 always visible
                    gr.update(visible=(num_runs >= 4)),  # output4
                    gr.update(visible=(num_runs >= 5))   # output5
                )
            
            generate_eval_btn.click(
                generate_three_versions,
                inputs=[eval_prompt, eval_temp1, eval_temp2, eval_temp3],
                outputs=[eval_output1, eval_output2, eval_output3]
            )
            
            calculate_ratings_btn.click(
                calculate_rating_summary,
                inputs=[
                    rate1_accuracy, rate1_coherence, rate1_completeness, rate1_relevance,
                    rate2_accuracy, rate2_coherence, rate2_completeness, rate2_relevance,
                    rate3_accuracy, rate3_coherence, rate3_completeness, rate3_relevance
                ],
                outputs=[ratings_summary]
            )
            
            generate_metric_btn.click(
                generate_and_compare,
                inputs=[metric_prompt, metric_reference, metric_temp],
                outputs=[metric_generated, metric_report, metric_overlap_display]
            )

    demo.launch()