Update app.py

reverting to old logic

app.py

@@ -3,6 +3,7 @@ import os
 import random
 import time
 import torch
+import re  # <-- NEW
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import spaces
 
@@ -23,18 +24,179 @@ except Exception as e:
     model = None
     tokenizer = None
 
-# --- Data for the Reels (Truncated for brevity) ---
+# --- 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: clean model output into a numbered list starting at "1." ---
+def clean_and_validate_output(text: str):
+    """
+    Extract the main numbered list starting at '1.' and strip any closing signature lines.
+    Always returns cleaned text and a boolean flag (True = looks fine).
+    """
+    # Keep everything from the first "1." onward, if present.
+    m = re.search(r'(?m)^\s*1\.\s', text)
+    body = text[m.start():].strip() if m else text.strip()
+
+    # Remove common signature lines at the end (best-effort).
+    body = re.sub(r'(?im)^\s*(yours.*|kind regards.*|regards.*)$', '', body).strip()
+
+    # If it doesn't contain at least one numbered point, it's still usable, but we mark as not strictly-valid.
+    is_valid = bool(re.search(r'(?m)^\s*\d+\.\s', body))
+    return body, is_valid
+
 # --- Helper: wrap content in the FOI letter template ---
 def wrap_in_letter(authority: str, body: str) -> str:
     body = body.strip()
@@ -46,15 +208,81 @@ def wrap_in_letter(authority: str, body: str) -> str:
     )
     return template
 
-# --- Gradio UI and Spinning Logic ---
-# --- Main Backend Function for Gradio ---
+# --- Backend Function for Local Inference ---
 @spaces.GPU
-def spin_and_generate():
-    # 1. Spin the reels for visual effect
-    spin_duration = 2.0
-    spin_interval = 0.05
+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 ONLY the numbered list of the specific information being requested, starting at '1.' "
+        f"for {authority}, using these keywords: {keyword_string}. "
+        "Do not include greetings or signatures."
+    )
+
+    max_retries = 2
+    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": 250,
+                    "do_sample": True,
+                    "temperature": 0.25,
+                    "top_k": 50,
+                    "top_p": 0.95,
+                    "repetition_penalty": 1.1,
+                    "streamer": None,
+                    "eos_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:]
+
+            # Clean and validate the output
+            cleaned_text, is_valid = clean_and_validate_output(generated_text)
+
+            # Wrap in the letter template regardless; validation just influences retry behavior
+            letter = wrap_in_letter(authority, cleaned_text)
+
+            if is_valid:
+                return letter
+            else:
+                print(f"Attempt {attempt + 1}/{max_retries}: Output lacked clear numbering. 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 retries failed, return the best effort letter using the last cleaned text we had
+    return wrap_in_letter(authority, "1. [Unable to format automatically] Please restate the information requested.\n2. [Optional second point]")
+
+# --- 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),
@@ -64,63 +292,30 @@ def spin_and_generate():
         )
         time.sleep(spin_interval)
 
-    # 2. Select the final combination and update UI
+    # 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))
-    kw1, kw2, kw3 = keywords_list[:3]
+    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. Generate the request using the model
-    if not model or not tokenizer:
-        yield (final_authority, kw1, kw2, kw3, "Error: Model is not loaded. Check Space logs.")
-        return
-
-    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 the request text (without greeting or signature) for {final_authority}, using these keywords: {keyword_string}."
-    )
-
-    try:
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-        generation_params = {
-            "max_new_tokens": 250,
-            "do_sample": True,
-            "temperature": 0.25,
-            "top_k": 50,
-            "top_p": 0.95,
-            "repetition_penalty": 1.1,
-            "eos_token_id": tokenizer.eos_token_id
-        }
-        output_sequences = model.generate(**inputs, **generation_params)
-        generated_text = tokenizer.decode(
-            output_sequences[0][len(inputs["input_ids"][0]):],
-            skip_special_tokens=True
-        ).strip()
-
-        if generated_text.startswith('.\n'):
-            generated_text = generated_text[2:]
-
-        final_request = wrap_in_letter(final_authority, generated_text)
-
-    except Exception as e:
-        final_request = f"An error occurred during text generation: {e}"
-
-    # 4. Yield the final result to the UI
+    # 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,
-        final_request
+        generated_request
     )
 
-
 # --- CSS for Styling ---
+# Added min-width to reduce UI flickering on text change
 reels_css = """
 #reels-container {
     display: flex;
@@ -133,7 +328,7 @@ reels_css = """
     border-radius: 12px;
     background-color: #fef3c7;
     box-shadow: 0 4px 6px rgba(0,0,0,0.1);
-    min-width: 150px;
+    min-width: 150px; /* Prevents resizing/flickering during spin */
 }
 #reels-container .gradio-textbox input {
     font-size: 1.25rem !important;
@@ -178,7 +373,7 @@ with gr.Blocks(css=reels_css, theme=gr.themes.Soft()) as demo:
     )
 
     pull_button.click(
-        fn=spin_and_generate, # Changed to the new combined function
+        fn=spin_the_reels,
         inputs=[],
         outputs=[reel1, reel2, reel3, reel4, output_request]
     )