app.py

import gradio as gr
import os
import random
import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import spaces

MODEL_ID = "HMC83/request_badger_olmo"

# --- Load Model and Tokenizer ---
print("Loading model and tokenizer...")
try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.bfloat16,
        device_map="auto"
    )
    print("Model loaded successfully.")
except Exception as e:
    print(f"Error loading model: {e}")
    model = None
    tokenizer = None

# --- Data for the Reels ---
# A list of authority and keyword combinations.
FOI_COMBINATIONS = [
    {"authority": "Borders NHS Board", "keywords": "whistleblowing guidance, wrongdoing, public body"},
    {"authority": "Borders NHS Board", "keywords": "ethical support, clinical triage, minutes"},
    {"authority": "Buckinghamshire Fire and Rescue Service", "keywords": "fire service, fire officers, membership"},
    {"authority": "Ballymoney Borough Council", "keywords": "parking enforcement, employee responsibility, parking services"},
    {"authority": "British Council", "keywords": "British Council funding, festival, correspondence"},
    {"authority": "Cardiff and Vale University Health Board", "keywords": "drugs, therapeutic, meetings"},
    {"authority": "Hastings Borough Council", "keywords": "local authority care, protection orders, adopted children"},
    {"authority": "Sevenoaks District Council", "keywords": "residential property, council tax, postcode"},
    {"authority": "The Marine Society and Sea Cadets", "keywords": "training, sea cadets, years in office"},
    {"authority": "West Yorkshire Police Authority", "keywords": "professional standards, official duties, chief constable"},
    {"authority": "Royal College of Veterinary Surgeons", "keywords": "veterinary practices, veterinary hospitals, contact details"},
    {"authority": "South Tyneside Metropolitan Borough Council", "keywords": "commercial customers, waste collection, waste producers"},
    {"authority": "University of Surrey", "keywords": "university car parks, parking tickets, parking enforcement"},
    {"authority": "King's College Hospital NHS Trust", "keywords": "annual spend, medical equipment, procurement"},
    {"authority": "Great Ormond Street Hospital", "keywords": "charity donations, research funding, financial year"},
    {"authority": "Public Health England", "keywords": "vaccination costs, public health campaigns, effectiveness"},
    {"authority": "NHS Digital", "keywords": "IT spending, system outages, cyber security"},
    {"authority": "Cornwall Council", "keywords": "adult social care, spending, assessments"},
    {"authority": "Hertfordshire County Council", "keywords": "school transport, route costs, contractor payments"},
    {"authority": "Wiltshire Council", "keywords": "highway maintenance, pothole repairs, annual spend"},
    {"authority": "Surrey County Council", "keywords": "council tax, collection rates, arrears"},
    {"authority": "Ofgem", "keywords": "energy company, penalties, investigations"},
    {"authority": "Financial Conduct Authority", "keywords": "enforcement action, fines imposed, complaint statistics"},
    {"authority": "Competition and Markets Authority", "keywords": "merger investigations, market studies, legal costs"},
    {"authority": "Health and Safety Executive", "keywords": "workplace accidents, enforcement notices, prosecution costs"},
    {"authority": "Network Rail", "keywords": "delay compensation, signal failures, maintenance costs"},
    {"authority": "Highways England", "keywords": "road maintenance, contractor spending, traffic flow"},
    {"authority": "Civil Aviation Authority", "keywords": "airline complaints, enforcement action, safety investigations"},
    {"authority": "HM Revenue and Customs", "keywords": "tax investigations, recovery rates, compliance costs"},
    {"authority": "Department for Work and Pensions", "keywords": "benefit sanctions, appeal outcomes, administrative costs"},
    {"authority": "Nottinghamshire County Council", "keywords": "adult social care, budget allocation, waiting lists"},
    {"authority": "Staffordshire County Council", "keywords": "school meals, contract costs, nutritional standards"},
    {"authority": "Derbyshire County Council", "keywords": "library services, closure consultations, usage statistics"},
    {"authority": "East Sussex County Council", "keywords": "children services, safeguarding, staffing costs"},
    {"authority": "Northumberland County Council", "keywords": "waste collection, recycling rates, contractor performance"},
    {"authority": "Royal Devon NHS Trust", "keywords": "patient complaints, response times, resolution outcomes"},
    {"authority": "Sheffield Teaching Hospitals", "keywords": "medical equipment, procurement costs, maintenance contracts"},
    {"authority": "Bradford Teaching Hospitals", "keywords": "staff training, external courses, annual expenditure"},
    {"authority": "Northampton General Hospital", "keywords": "parking charges, revenue generated, patient feedback"},
    {"authority": "Environment Agency", "keywords": "flood defenses, maintenance costs, effectiveness assessments"},
    {"authority": "Natural England", "keywords": "habitat protection, designation costs, landowner compensation"},
    {"authority": "Marine Management Organisation", "keywords": "fishing licenses, enforcement costs, violation penalties"},
    {"authority": "Forestry Commission", "keywords": "timber sales, revenue targets, environmental impact"},
    {"authority": "Welsh Government", "keywords": "language services, translation costs, usage statistics"},
    {"authority": "Scottish Government", "keywords": "ferry subsidies, route profitability, passenger numbers"},
    {"authority": "Northern Ireland Executive", "keywords": "cross-border initiatives, funding allocation, outcomes"},
    {"authority": "Office for National Statistics", "keywords": "census costs, data collection, contractor payments"},
    {"authority": "Land Registry", "keywords": "property registrations, processing times, fee income"},
    {"authority": "Companies House", "keywords": "late filing penalties, collection rates, enforcement action"},
    {"authority": "Intellectual Property Office", "keywords": "patent applications, processing costs, appeal outcomes"},
    {"authority": "InGen Safety Authority", "keywords": "containment breaches, insurance costs, visitor incidents"},
    {"authority": "Natural England", "keywords": "Bigfoot sightings, habitat, evidence"},
    {"authority": "The Royal Parks", "keywords": "squirrel uprising, acorn stockpiles, intelligence"},
    {"authority": "Midsomer Constabulary", "keywords": "unusually high murder rate, statistics, investigations"},
    {"authority": "Scottish Government", "keywords": "Loch Ness Monster, research, funding"},
    {"authority": "Westminster City Council", "keywords": "noise complaints, nightlife, enforcement"},
    {"authority": "Network Rail", "keywords": "delays, compensation, signal failures"},
    {"authority": "Gloucestershire County Council", "keywords": "adult social care, assessment costs, service provision"},
    {"authority": "Worcestershire County Council", "keywords": "school transport, route planning, contractor performance"},
    {"authority": "Leicestershire County Council", "keywords": "highway maintenance, annual spend, emergency repairs"},
    {"authority": "Warwickshire County Council", "keywords": "library services, usage statistics, closure consultations"},
    {"authority": "Nottinghamshire County Council", "keywords": "children services, safeguarding referrals, staff costs"},
    {"authority": "Derbyshire County Council", "keywords": "waste disposal, recycling rates, collection contracts"},
    {"authority": "Staffordshire County Council", "keywords": "care homes, inspection reports, quality ratings"},
    {"authority": "Shropshire Council", "keywords": "planning applications, approval rates, enforcement action"},
    {"authority": "Herefordshire Council", "keywords": "council tax, collection rates, arrears management"},
    {"authority": "Cheshire East Council", "keywords": "parking enforcement, penalty notices, appeal outcomes"},
    {"authority": "Royal Liverpool University Hospital", "keywords": "patient complaints, response times, resolution outcomes"},
    {"authority": "Sheffield Teaching Hospitals NHS Trust", "keywords": "medical equipment, procurement costs, maintenance contracts"},
    {"authority": "Leeds Teaching Hospitals NHS Trust", "keywords": "agency staff, temporary costs, recruitment spend"},
    {"authority": "Newcastle upon Tyne Hospitals NHS Trust", "keywords": "cancelled operations, reasons given, patient impact"},
    {"authority": "University Hospitals Birmingham", "keywords": "waiting times, treatment delays, performance targets"},
    {"authority": "Guy's and St Thomas' NHS Trust", "keywords": "infection control, outbreak incidents, prevention costs"},
    {"authority": "King's College Hospital NHS Trust", "keywords": "emergency admissions, capacity management, staffing levels"},
    {"authority": "Royal Free NHS Trust", "keywords": "clinical trials, research funding, patient consent"},
    {"authority": "Imperial College Healthcare NHS Trust", "keywords": "private patient, income generated, capacity allocation"},
    {"authority": "Barts Health NHS Trust", "keywords": "financial deficit, recovery plans, cost reduction"},
    {"authority": "West Yorkshire Police", "keywords": "response times, emergency calls, performance standards"},
    {"authority": "South Yorkshire Police", "keywords": "crime statistics, detection rates, resource allocation"},
    {"authority": "North Yorkshire Police", "keywords": "rural policing, coverage areas, response costs"},
    {"authority": "Humberside Police", "keywords": "traffic enforcement, speed cameras, revenue generated"},
    {"authority": "Lincolnshire Police", "keywords": "domestic violence, incident reports, support services"},
    {"authority": "Nottinghamshire Police", "keywords": "stop and search, statistics recorded, outcome data"},
    {"authority": "Derbyshire Police", "keywords": "drug seizures, investigation costs, conviction rates"},
    {"authority": "Leicestershire Police", "keywords": "neighbourhood policing, community engagement, effectiveness measures"},
    {"authority": "Warwickshire Police", "keywords": "police stations, closure plans, public consultations"},
    {"authority": "West Mercia Police", "keywords": "cybercrime, investigation resources, training costs"},
    {"authority": "Environment Agency", "keywords": "flood defenses, maintenance spending, effectiveness assessments"},
    {"authority": "Natural England", "keywords": "habitat designations, compensation payments, landowner agreements"},
    {"authority": "Historic England", "keywords": "listed buildings, consent applications, enforcement cases"},
    {"authority": "Highways England", "keywords": "motorway maintenance, contractor payments, performance indicators"},
    {"authority": "Network Rail", "keywords": "signal failures, delay minutes, compensation costs"},
    {"authority": "Civil Aviation Authority", "keywords": "airline complaints, resolution times, enforcement action"},
    {"authority": "Maritime and Coastguard Agency", "keywords": "rescue operations, helicopter costs, response times"},
    {"authority": "Driver and Vehicle Standards Agency", "keywords": "MOT failures, test center, compliance rates"},
    {"authority": "Vehicle Certification Agency", "keywords": "emissions testing, approval procedures, manufacturer compliance"},
    {"authority": "Planning Inspectorate", "keywords": "appeal decisions, processing times, outcome statistics"},
    {"authority": "Health and Safety Executive", "keywords": "workplace accidents, investigation costs, prosecution outcomes"},
    {"authority": "Care Quality Commission", "keywords": "inspection reports, rating changes, enforcement action"},
    {"authority": "Ofgem", "keywords": "energy company, penalties imposed, consumer complaints"},
    {"authority": "Ofwat", "keywords": "water companies, price reviews, performance monitoring"},
    {"authority": "Ofcom", "keywords": "broadcasting complaints, investigation procedures, penalty decisions"},
    {"authority": "Financial Conduct Authority", "keywords": "financial services, enforcement cases, penalty amounts"},
    {"authority": "Competition and Markets Authority", "keywords": "merger investigations, market studies, intervention costs"},
    {"authority": "Gambling Commission", "keywords": "license suspensions, operator penalties, compliance monitoring"},
    {"authority": "Information Commissioner's Office", "keywords": "data breaches, penalty notices, investigation outcomes"},
    {"authority": "Electoral Commission", "keywords": "campaign spending, compliance checks, penalty procedures"},
    {"authority": "HM Revenue and Customs", "keywords": "tax investigations, recovery amounts, compliance costs"},
    {"authority": "HM Treasury", "keywords": "departmental budgets, spending reviews, efficiency targets"},
    {"authority": "Cabinet Office", "keywords": "government consultancy, external advisors, procurement costs"},
    {"authority": "Ministry of Justice", "keywords": "court delays, case backlogs, administrative costs"},
    {"authority": "Home Office", "keywords": "immigration appeals, processing times, detention costs"},
    {"authority": "Department for Transport", "keywords": "railway subsidies, franchise performance, punctuality targets"},
    {"authority": "Department for Education", "keywords": "academy conversions, funding allocations, performance data"},
    {"authority": "Department for Work and Pensions", "keywords": "benefit sanctions, appeal success, administrative costs"},
    {"authority": "Department of Health", "keywords": "NHS funding, allocation formulas, performance targets"},
    {"authority": "Department for Environment", "keywords": "environmental fines, pollution incidents, enforcement costs"},
    {"authority": "London Fire Brigade", "keywords": "response times, equipment costs, staffing levels"},
    {"authority": "West Midlands Fire Service", "keywords": "false alarms, call reduction, prevention costs"},
    {"authority": "Greater Manchester Fire Service", "keywords": "building inspections, safety compliance, enforcement action"},
    {"authority": "South Yorkshire Fire Service", "keywords": "emergency calls, resource deployment, performance metrics"},
    {"authority": "Merseyside Fire Service", "keywords": "community safety, education programs, effectiveness measures"},
    {"authority": "Essex Fire Service", "keywords": "rescue operations, specialist equipment, training costs"},
    {"authority": "Kent Fire Service", "keywords": "station closures, service changes, public consultation"},
    {"authority": "Surrey Fire Service", "keywords": "road accidents, rescue costs, casualty statistics"},
    {"authority": "Hampshire Fire Service", "keywords": "wildfire incidents, prevention measures, resource allocation"},
    {"authority": "Devon Fire Service", "keywords": "coastal rescues, equipment maintenance, operational costs"},
    {"authority": "London Ambulance Service", "keywords": "response targets, performance data, resource pressures"},
    {"authority": "West Midlands Ambulance", "keywords": "call volumes, crew availability, overtime costs"},
    {"authority": "Yorkshire Ambulance Service", "keywords": "hospital handovers, delay times, capacity issues"},
    {"authority": "North West Ambulance Service", "keywords": "emergency calls, triage procedures, response priorities"},
    {"authority": "South Central Ambulance", "keywords": "patient transport, contract costs, service quality"},
    {"authority": "East Midlands Ambulance", "keywords": "staff training, competency assessments, certification costs"},
    {"authority": "South East Coast Ambulance", "keywords": "vehicle maintenance, fleet costs, replacement schedules"},
    {"authority": "East of England Ambulance", "keywords": "clinical governance, audit results, quality improvements"},
    {"authority": "South Western Ambulance", "keywords": "rural coverage, response challenges, resource deployment"},
    {"authority": "Welsh Ambulance Service", "keywords": "cross-border operations, coordination costs, service agreements"},
    {"authority": "Bristol City Council", "keywords": "cycling infrastructure, usage data, maintenance costs"},
    {"authority": "Brighton and Hove Council", "keywords": "seafront management, event licensing, revenue generation"},
    {"authority": "Cambridge City Council", "keywords": "student housing, planning enforcement, compliance monitoring"},
    {"authority": "Canterbury City Council", "keywords": "heritage sites, conservation costs, visitor management"},
    {"authority": "Chester West Council", "keywords": "tourism promotion, marketing spend, economic impact"},
    {"authority": "Colchester Borough Council", "keywords": "business rates, collection performance, appeals process"},
    {"authority": "Durham County Council", "keywords": "mining subsidence, compensation claims, repair costs"},
    {"authority": "Exeter City Council", "keywords": "flood management, defense systems, emergency planning"},
    {"authority": "Gloucester City Council", "keywords": "regeneration projects, funding sources, completion rates"},
    {"authority": "Lancaster City Council", "keywords": "coastal erosion, protection measures, maintenance spending"},
]

# Create lists for the spinning animation from the combinations above
ALL_AUTHORITIES_FOR_SPIN = list(set([item["authority"] for item in FOI_COMBINATIONS]))
ALL_KEYWORDS_FOR_SPIN = list(set(kw.strip() for item in FOI_COMBINATIONS for kw in item["keywords"].split(',')))


# --- Helper Function for Cleaning and Validation ---
def clean_and_validate_output(raw_text: str) -> tuple[str, bool]:
    """
    Cleans the model's output by keeping only the first complete request.

    It validates that the output contains essential markers ("Dear" and "[Your Name]").
    If it detects that the model has started generating a second request, it truncates
    the string after the first "[Your Name]".

    Args:
        raw_text: The raw string output from the language model.

    Returns:
        A tuple containing:
        - The cleaned text.
        - A boolean flag: True if the output is valid, False if it is malformed.
    """
    end_marker = "[Your Name]"
    start_marker = "Dear"

    # Validate: A valid request must contain the end marker.
    if end_marker not in raw_text:
        return raw_text, False  # Malformed, signal for regeneration.

    # Find the end of the first complete request.
    first_end_pos = raw_text.find(end_marker)
    end_of_first_request_index = first_end_pos + len(end_marker)

    # Check if a second request has started after the first one ended.
    start_of_second_request_pos = raw_text.find(start_marker, end_of_first_request_index)

    if start_of_second_request_pos != -1:
        # If a second request is found, truncate to keep only the first one.
        cleaned_text = raw_text[:end_of_first_request_index]
        return cleaned_text, True
    else:
        # No second request found, the output is valid.
        return raw_text, True


# --- Backend Function for Local Inference ---
@spaces.GPU
def generate_request_local(authority, kw1, kw2, kw3):
    """Generates a request using the locally loaded transformer model, with validation and retry logic."""
    if not model or not tokenizer:
        return "Error: Model is not loaded. Please check the Space logs for details."

    keywords = [kw for kw in [kw1, kw2, kw3] if kw]
    keyword_string = ", ".join(keywords)
    prompt = (
        "You are an expert at writing formal Freedom of Information requests to UK public authorities. "
        f"""Generate a formal Freedom of Information request to {authority} using these keywords: {keyword_string}
             The request should:
             1. Start with "Dear {authority},"
             2. Continue with "please provide me with a copy of the following information:"
             3. Include a numbered list of 1-2 requests for specific information related to the keywords
             4. The requests should not be too broad, and should be ones that a typical journalist or member of the public might make.
             5. The requests must not ask for publicly available information
             6. They must clearly describe the exact information being requested, in a way that doesn't require any subjective judgement
             on the part of the FOI officer to answer the request. No clarification should be needed.
             7. Do not ask for all policies, or all information
             8. End with "Yours Faithfully, [Your Name]" exactly
            Make the requests specific, professional, and relevant to what this public authority would reasonably hold.
            Use accessible language, avoiding terms that are overly legalistic or technical and UK English. Be clear and concise"""
    )

    max_retries = 3
    for attempt in range(max_retries):
        try:
            # Tokenize the input prompt
            inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

            # Set generation parameters
            generation_params = {
                "max_new_tokens": 500,
                "temperature": 0.3,
                "top_p": 0.95,
                "top_k": 50,
                "repetition_penalty": 1.05,
                "do_sample": True,
                "pad_token_id": tokenizer.eos_token_id
            }

            # Generate text sequences
            output_sequences = model.generate(**inputs, **generation_params)

            # Decode the generated text
            generated_text = tokenizer.decode(
                output_sequences[0][len(inputs["input_ids"][0]):],
                skip_special_tokens=True
            ).strip()

            # Remove artifact if present
            if generated_text.startswith('.\n'):
                generated_text = generated_text[2:]

            # **NEW**: Clean and validate the output
            cleaned_text, is_valid = clean_and_validate_output(generated_text)
            
            if is_valid:
                return cleaned_text  # Success! Return the valid, cleaned text.
            else:
                print(f"Attempt {attempt + 1}/{max_retries}: Malformed output detected. Retrying...")

        except Exception as e:
            print(f"Error during generation attempt {attempt + 1}/{max_retries}: {e}")
            if attempt == max_retries - 1:
                return f"An error occurred during text generation: {e}"

    # If the loop finishes, all retries have failed
    return "Failed to generate a valid request after multiple attempts. Please try again."


# --- Gradio UI and Spinning Logic ---
def spin_the_reels():
    """A generator function that simulates spinning reels and then calls the model."""
    # 1. Simulate the spinning effect
    spin_duration = 2.0  # seconds
    spin_interval = 0.05 # update interval
    start_time = time.time()
    while time.time() - start_time < spin_duration:
        # Yield random values for each reel to create the spinning illusion
        yield (
            random.choice(ALL_AUTHORITIES_FOR_SPIN),
            random.choice(ALL_KEYWORDS_FOR_SPIN),
            random.choice(ALL_KEYWORDS_FOR_SPIN),
            random.choice(ALL_KEYWORDS_FOR_SPIN),
            "Spinning..."
        )
        time.sleep(spin_interval)
    
    # 2. Select the final fixed combination
    final_combination = random.choice(FOI_COMBINATIONS)
    final_authority = final_combination["authority"]
    
    # Split, strip, and pad keywords to ensure we always have 3 for the UI
    keywords_list = [k.strip() for k in final_combination["keywords"].split(',')]
    keywords_list += [''] * (3 - len(keywords_list)) # Pad with empty strings if < 3
    kw1, kw2, kw3 = keywords_list[:3] # Take the first 3
    
    # Display the final reel values and a "Generating..." message
    yield (
        final_authority, kw1, kw2, kw3,
        f"Generating request for {final_authority}...\nPlease wait, this may take a moment."
    )
    
    # 3. Call the local model and yield the final result
    generated_request = generate_request_local(final_authority, kw1, kw2, kw3)
    yield (
        final_authority, kw1, kw2, kw3,
        generated_request
    )

# --- CSS for Styling ---
# Added min-width to reduce UI flickering on text change
reels_css = """
#reels-container {
    display: flex;
    gap: 1rem;
    margin-bottom: 1rem;
}
#reels-container .gradio-textbox {
    text-align: center;
    border: 4px solid #f59e0b;
    border-radius: 12px;
    background-color: #fef3c7;
    box-shadow: 0 4px 6px rgba(0,0,0,0.1);
    min-width: 150px; /* Prevents resizing/flickering during spin */
}
#reels-container .gradio-textbox input {
    font-size: 1.25rem !important;
    font-weight: bold;
    color: #b45309;
    text-align: center;
}
#pull-button {
    font-size: 1.5rem !important;
    font-weight: bold;
    padding: 1rem !important;
    box-shadow: 0 8px 15px rgba(0,0,0,0.2);
    transition: all 0.2s ease-in-out;
}
#pull-button:hover {
    transform: translateY(-2px);
    box-shadow: 0 10px 20px rgba(0,0,0,0.3);
}
"""

# --- Build the Gradio App ---
with gr.Blocks(css=reels_css, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # FOI Request-O-Matic 🎰
        Press the button to generate a Freedom of Information request.
        """
    )
    with gr.Row(elem_id="reels-container"):
        reel1 = gr.Textbox(label="Authority", interactive=False, elem_id="reel-1", scale=2)
        reel2 = gr.Textbox(label="Keyword 1", interactive=False, elem_id="reel-2", scale=1)
        reel3 = gr.Textbox(label="Keyword 2", interactive=False, elem_id="reel-3", scale=1)
        reel4 = gr.Textbox(label="Keyword 3", interactive=False, elem_id="reel-4", scale=1)
    
    pull_button = gr.Button("Generate a request", variant="primary", elem_id="pull-button")
    
    output_request = gr.Textbox(
        label="Generated FOI Request",
        lines=15,
        interactive=True,
        show_copy_button=True
    )

    pull_button.click(
        fn=spin_the_reels,
        inputs=[],
        outputs=[reel1, reel2, reel3, reel4, output_request]
    )

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