Update app.py

better handling of messier outputs

app.py

@@ -180,16 +180,55 @@ ALL_AUTHORITIES_FOR_SPIN = list(set([item["authority"] for item in FOI_COMBINATI
 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 on a dynamically allocated GPU."""
+    """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}
@@ -203,48 +242,56 @@ def generate_request_local(authority, kw1, kw2, kw3):
              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"""
     )
-        
-    try:
-        # Tokenize the input prompt
-        inputs = tokenizer(prompt, return_tensors="pt")
-        
-        # Move to the model's device (should be cuda when GPU is allocated)
-        inputs = inputs.to(model.device)
-        
-        # Set generation parameters to match notebook for better performance
-        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, skipping special tokens and the original prompt
-        generated_text = tokenizer.decode(
-            output_sequences[0][len(inputs["input_ids"][0]):],
-            skip_special_tokens=True
-        ).strip()
-        
-        # Add this check to remove artifact
-        if generated_text.startswith('.\n'):
-            generated_text = generated_text[2:] 
-        return generated_text
-        
-    except Exception as e:
-        print(f"Error during generation: {e}")
-        print(f"Model device: {model.device if model else 'No model'}")
-        print(f"Inputs type: {type(inputs) if 'inputs' in locals() else 'Not created'}")
-        return f"An error occurred during text generation: {e}"
+
+    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():
@@ -263,7 +310,7 @@ def spin_the_reels():
             "Spinning..."
         )
         time.sleep(spin_interval)
-
+    
     # 2. Select the final fixed combination
     final_combination = random.choice(FOI_COMBINATIONS)
     final_authority = final_combination["authority"]
@@ -278,7 +325,7 @@ def spin_the_reels():
         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 (
@@ -334,7 +381,7 @@ with gr.Blocks(css=reels_css, theme=gr.themes.Soft()) as demo:
         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(