Update app.py (#3)

- Update app.py (f46c493ee0918191301804380b8116814c315750)

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import logging
 
-# Setup logging (optional, but helpful for debugging)
+# Setup logging
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # Load the Flan-T5 Small model and tokenizer
@@ -12,81 +12,54 @@ model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
 
 def correct_htr(raw_htr_text, max_new_tokens, temperature):
     try:
-        if not raw_htr_text:
-            raise ValueError("Input text cannot be empty.")
-        
-        logging.info("Processing HTR correction with Flan-T5 Small...")
+        logging.info("Processing HTR correction...")
         prompt = f"Correct this text: {raw_htr_text}"
         inputs = tokenizer(prompt, return_tensors="pt")
-        max_length = min(max_new_tokens, len(inputs['input_ids'][0]) + max_new_tokens)
-        outputs = model.generate(**inputs, max_length=max_length, temperature=temperature)
+        outputs = model.generate(**inputs, max_length=min(max_new_tokens, len(inputs['input_ids'][0]) + max_new_tokens), temperature=temperature)
         corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        logging.debug(f"Generated output for HTR correction: {corrected_text}")
         return corrected_text
-    except ValueError as ve:
-        logging.warning(f"Validation error: {ve}")
-        return str(ve)
     except Exception as e:
         logging.error(f"Error in HTR correction: {e}", exc_info=True)
-        return "An error occurred while processing the text."
+        return str(e)
 
 def summarize_text(legal_text, max_new_tokens, temperature):
     try:
-        if not legal_text:
-            raise ValueError("Input text cannot be empty.")
-        
-        logging.info("Processing summarization with Flan-T5 Small...")
+        logging.info("Processing summarization...")
         prompt = f"Summarize the following legal text: {legal_text}"
         inputs = tokenizer(prompt, return_tensors="pt")
-        max_length = min(max_new_tokens, len(inputs['input_ids'][0]) + max_new_tokens)
-        outputs = model.generate(**inputs, max_length=max_length, temperature=temperature)
+        outputs = model.generate(**inputs, max_length=min(max_new_tokens, len(inputs['input_ids'][0]) + max_new_tokens), temperature=temperature)
         summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        logging.debug(f"Generated summary: {summary}")
         return summary
-    except ValueError as ve:
-        logging.warning(f"Validation error: {ve}")
-        return str(ve)
     except Exception as e:
         logging.error(f"Error in summarization: {e}", exc_info=True)
-        return "An error occurred while summarizing the text."
+        return str(e)
 
 def answer_question(legal_text, question, max_new_tokens, temperature):
     try:
-        if not legal_text or not question:
-            raise ValueError("Both legal text and question must be provided.")
-        
-        logging.info("Processing question-answering with Flan-T5 Small...")
+        logging.info("Processing question-answering...")
         prompt = f"Answer the following question based on the provided context:\n\nQuestion: {question}\n\nContext: {legal_text}"
         inputs = tokenizer(prompt, return_tensors="pt")
-        max_length = min(max_new_tokens, len(inputs['input_ids'][0]) + max_new_tokens)
-        outputs = model.generate(**inputs, max_length=max_length, temperature=temperature)
+        outputs = model.generate(**inputs, max_length=min(max_new_tokens, len(inputs['input_ids'][0]) + max_new_tokens), temperature=temperature)
         answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        logging.debug(f"Generated answer: {answer}")
         return answer
-    except ValueError as ve:
-        logging.warning(f"Validation error: {ve}")
-        return str(ve)
     except Exception as e:
         logging.error(f"Error in question-answering: {e}", exc_info=True)
-        return "An error occurred while answering the question."
-
-def clear_fields():
-    return "", "", ""
+        return str(e)
 
 # Create the Gradio Blocks interface
-with gr.Blocks(css=".block .input-slider { color: blue !important }") as demo:
+with gr.Blocks() as demo:
     gr.Markdown("# Flan-T5 Small Legal Assistant")
     gr.Markdown("Use this tool to correct raw HTR, summarize legal texts, or answer questions about legal cases (powered by Flan-T5 Small).")
-
+    
     with gr.Row():
         gr.HTML('''
             <div style="display: flex; gap: 10px;">
-                <div style="border: 2px solid black; padding: 10px; display: inline-block;">
+                <div style="border: 2px solid black; padding: 10px;">
                     <a href="http://www.marinelives.org/wiki/Tools:_Admiralty_court_legal_glossary" target="_blank">
                         <button style="font-weight:bold;">Admiralty Court Legal Glossary</button>
                     </a>
                 </div>
-                <div style="border: 2px solid black; padding: 10px; display: inline-block;">
+                <div style="border: 2px solid black; padding: 10px;">
                     <a href="https://raw.githubusercontent.com/Addaci/HCA/refs/heads/main/HCA_13_70_Full_Volume_Processed_Text_EDITED_Ver.1.2_18062024.txt" target="_blank">
                         <button style="font-weight:bold;">HCA 13/70 Ground Truth (1654-55)</button>
                     </a>
@@ -94,38 +67,27 @@ with gr.Blocks(css=".block .input-slider { color: blue !important }") as demo:
             </div>
         ''')
 
+    # Tab 1: Correct HTR
     with gr.Tab("Correct HTR"):
         gr.Markdown("### Correct Raw HTR Text")
         raw_htr_input = gr.Textbox(lines=5, placeholder="Enter raw HTR text here...")
         corrected_output = gr.Textbox(lines=5, placeholder="Corrected HTR text")
         correct_button = gr.Button("Correct HTR")
         clear_button = gr.Button("Clear")
-        correct_max_new_tokens = gr.Slider(minimum=10, maximum=512, value=128, step=1, label="Max New Tokens")
-        correct_temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
-
-        correct_button.click(correct_htr, inputs=[raw_htr_input, correct_max_new_tokens, correct_temperature], outputs=corrected_output)
-        clear_button.click(clear_fields, outputs=[raw_htr_input, corrected_output])
-        
-        gr.Markdown("### Set Parameters")
-        correct_max_new_tokens.render()
-        correct_temperature.render()
+        correct_button.click(correct_htr, inputs=[raw_htr_input, gr.Slider(minimum=10, maximum=512, value=128, step=1, label="Max New Tokens"), gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")], outputs=corrected_output)
+        clear_button.click(lambda: ("", ""), outputs=[raw_htr_input, corrected_output])
 
+    # Tab 2: Summarize Legal Text
     with gr.Tab("Summarize Legal Text"):
         gr.Markdown("### Summarize Legal Text")
         legal_text_input = gr.Textbox(lines=10, placeholder="Enter legal text to summarize...")
         summary_output = gr.Textbox(lines=5, placeholder="Summary of legal text")
         summarize_button = gr.Button("Summarize Text")
         clear_button = gr.Button("Clear")
-        summarize_max_new_tokens = gr.Slider(minimum=10, maximum=1024, value=256, step=1, label="Max New Tokens")
-        summarize_temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.1, label="Temperature")
-
-        summarize_button.click(summarize_text, inputs=[legal_text_input, summarize_max_new_tokens, summarize_temperature], outputs=summary_output)
-        clear_button.click(clear_fields, outputs=[legal_text_input, summary_output])
-        
-        gr.Markdown("### Set Parameters")
-        summarize_max_new_tokens.render()
-        summarize_temperature.render()
+        summarize_button.click(summarize_text, inputs=[legal_text_input, gr.Slider(minimum=10, maximum=512, value=256, step=1, label="Max New Tokens"), gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.1, label="Temperature")], outputs=summary_output)
+        clear_button.click(lambda: ("", ""), outputs=[legal_text_input, summary_output])
 
+    # Tab 3: Answer Legal Question
     with gr.Tab("Answer Legal Question"):
         gr.Markdown("### Answer a Question Based on Legal Text")
         legal_text_input_q = gr.Textbox(lines=10, placeholder="Enter legal text...")
@@ -133,18 +95,8 @@ with gr.Blocks(css=".block .input-slider { color: blue !important }") as demo:
         answer_output = gr.Textbox(lines=5, placeholder="Answer to your question")
         answer_button = gr.Button("Get Answer")
         clear_button = gr.Button("Clear")
-        answer_max_new_tokens = gr.Slider(minimum=10, maximum=512, value=150, step=1, label="Max New Tokens")
-        answer_temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, label="Temperature")
-
-        answer_button.click(answer_question, inputs=[legal_text_input_q, question_input, answer_max_new_tokens, answer_temperature], outputs=answer_output)
-        clear_button.click(clear_fields, outputs=[legal_text_input_q, question_input, answer_output])
-        
-        gr.Markdown("### Set Parameters")
-        answer_max_new_tokens.render()
-        answer_temperature.render()
-
-# Model warm-up (optional, but useful for performance)
-model.generate(**tokenizer("Warm-up", return_tensors="pt"), max_length=10)
+        answer_button.click(answer_question, inputs=[legal_text_input_q, question_input, gr.Slider(minimum=10, maximum=512, value=150, step=1, label="Max New Tokens"), gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, label="Temperature")], outputs=answer_output)
+        clear_button.click(lambda: ("", "", ""), outputs=[legal_text_input_q, question_input, answer_output])
 
 # Launch the Gradio interface
 if __name__ == "__main__":