Major rewrite of app.py

Rewrote the app.py file to apply a three tab, two button template from a different Hugging Face space
Added slide baars at bottom for text output size and temperature, with base values, and minimum and maxim,um values
Changed model to google/flan-t5-xl
This model, according to Gemini Pro: "is known for its strong performance in a wide range of tasks, including summarization and question answering"

Gemini also suggested looking at:

BART (facebook/bart-large-cnn): BART is another excellent choice, especially for summarization tasks.
LongT5 (google/long-t5-tglobal-xl): If you're dealing with longer legal texts, LongT5 might be a good option due to its ability to handle longer input sequences.

app.py

@@ -1,53 +1,58 @@
-# Cell 1B: Inference Client
-
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import logging
 
-# Setup logging
+# Setup logging (optional, but helpful for debugging)
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Initialize Inference Client
-client = InferenceClient(model="Addaci/mT5-small-experiment-13-checkpoint-2790")
+# Load the Flan-T5 model and tokenizer
+model_id = "google/flan-t5-xl"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
 
-def correct_htr(raw_htr_text):
+def correct_htr(raw_htr_text, max_new_tokens, temperature):
     try:
-        logging.info("Processing HTR correction with InferenceClient...")
-        # Sending the input to the hosted model
-        result = client.text_generation(f"correct this text: {raw_htr_text}")
-        logging.debug(f"Generated output for HTR correction: {result}")
-        return result['generated_text']  # Extracting the generated text from the response
+        logging.info("Processing HTR correction with Flan-T5...")
+        prompt = f"Correct this text: {raw_htr_text}"
+        inputs = tokenizer(prompt, return_tensors="pt")
+        outputs = model.generate(**inputs, max_length=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 Exception as e:
         logging.error(f"Error in HTR correction: {e}", exc_info=True)
         return str(e)
 
-def summarize_text(legal_text):
+def summarize_text(legal_text, max_new_tokens, temperature):
     try:
-        logging.info("Processing summarization with InferenceClient...")
-        # Sending the input to the hosted model
-        result = client.text_generation(f"summarize the following legal text: {legal_text}")
-        logging.debug(f"Generated summary: {result}")
-        return result['generated_text']  # Extracting the generated text from the response
+        logging.info("Processing summarization with Flan-T5...")
+        prompt = f"Summarize the following legal text: {legal_text}"
+        inputs = tokenizer(prompt, return_tensors="pt")
+        outputs = model.generate(**inputs, max_length=max_new_tokens, temperature=temperature)
+        summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        logging.debug(f"Generated summary: {summary}")
+        return summary
     except Exception as e:
         logging.error(f"Error in summarization: {e}", exc_info=True)
         return str(e)
 
-def answer_question(legal_text, question):
+def answer_question(legal_text, question, max_new_tokens, temperature):
     try:
-        logging.info("Processing question-answering with InferenceClient...")
-        # Sending the input to the hosted model
-        formatted_input = f"Answer the following question based on the provided context:\n\nQuestion: {question}\n\nContext: {legal_text}"
-        result = client.text_generation(formatted_input)
-        logging.debug(f"Generated answer: {result}")
-        return result['generated_text']  # Extracting the generated text from the response
+        logging.info("Processing question-answering with Flan-T5...")
+        prompt = f"Answer the following question based on the provided context:\n\nQuestion: {question}\n\nContext: {legal_text}"
+        inputs = tokenizer(prompt, return_tensors="pt")
+        outputs = model.generate(**inputs, max_length=max_new_tokens, temperature=temperature)
+        answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        logging.debug(f"Generated answer: {answer}")
+        return answer
     except Exception as e:
         logging.error(f"Error in question-answering: {e}", exc_info=True)
         return str(e)
 
 # Create the Gradio Blocks interface
 with gr.Blocks() as demo:
-    gr.Markdown("# mT5 Legal Assistant")
-    gr.Markdown("Use this tool to correct raw HTR, summarize legal texts, or answer questions about legal cases.")
+    gr.Markdown("# Flan-T5 Legal Assistant")
+    gr.Markdown("Use this tool to correct raw HTR, summarize legal texts, or answer questions about legal cases (powered by Flan-T5).")
 
     with gr.Row():
         gr.HTML('''
@@ -71,8 +76,8 @@ with gr.Blocks() as demo:
         corrected_output = gr.Textbox(lines=5, placeholder="Corrected HTR text")
         correct_button = gr.Button("Correct HTR")
         clear_button = gr.Button("Clear")
-        
-        correct_button.click(correct_htr, inputs=raw_htr_input, outputs=corrected_output)
+
+        correct_button.click(correct_htr, inputs=[raw_htr_input, max_new_tokens, temperature], outputs=corrected_output)
         clear_button.click(lambda: ("", ""), outputs=[raw_htr_input, corrected_output])
 
     with gr.Tab("Summarize Legal Text"):
@@ -81,8 +86,8 @@ with gr.Blocks() as demo:
         summary_output = gr.Textbox(lines=5, placeholder="Summary of legal text")
         summarize_button = gr.Button("Summarize Text")
         clear_button = gr.Button("Clear")
-        
-        summarize_button.click(summarize_text, inputs=legal_text_input, outputs=summary_output)
+
+        summarize_button.click(summarize_text, inputs=[legal_text_input, max_new_tokens, temperature], outputs=summary_output)
         clear_button.click(lambda: ("", ""), outputs=[legal_text_input, summary_output])
 
     with gr.Tab("Answer Legal Question"):
@@ -92,9 +97,14 @@ with gr.Blocks() 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_button.click(answer_question, inputs=[legal_text_input_q, question_input], outputs=answer_output)
+
+        answer_button.click(answer_question, inputs=[legal_text_input_q, question_input, max_new_tokens, temperature], outputs=answer_output)
         clear_button.click(lambda: ("", "", ""), outputs=[legal_text_input_q, question_input, answer_output])
 
+    # Add sliders for hyperparameters
+    with gr.Row():
+        max_new_tokens = gr.Slider(minimum=10, maximum=1000, value=500, step=1, label="Max New Tokens")
+        temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
+
 # Launch the Gradio interface
 demo.launch()