Update app.py (reverted to MT%Tokenizer; added further logging function)

app.py

@@ -1,15 +1,15 @@
 import os
 import gradio as gr
 import logging
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import MT5Tokenizer, MT5ForConditionalGeneration
 
 # Setup logging
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # Load your fine-tuned mT5 model
 model_name = "Addaci/mT5-small-experiment-13-checkpoint-2790"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+tokenizer = MT5Tokenizer.from_pretrained(model_name)
+model = MT5ForConditionalGeneration.from_pretrained(model_name)
 
 def correct_htr(raw_htr_text):
     try:
@@ -17,12 +17,16 @@ def correct_htr(raw_htr_text):
         inputs = tokenizer("correct this text: " + raw_htr_text, return_tensors="pt", max_length=512, truncation=True)
         logging.debug(f"Tokenized Inputs for HTR Correction: {inputs}")
 
-        # Set do_sample=True
-        outputs = model.generate(**inputs, max_length=128, num_beams=4, early_stopping=True, temperature=0.6, do_sample=True)  
+        # Generate with beam search and sampling
+        outputs = model.generate(**inputs, max_length=128, num_beams=4, early_stopping=True, temperature=0.6, do_sample=True)
         logging.debug(f"Generated Output (Tokens) for HTR Correction: {outputs}")
 
         corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         logging.debug(f"Decoded Output for HTR Correction: {corrected_text}")
+
+        # Re-tokenize the output for further inspection
+        logging.debug(f"Re-tokenized output for HTR Correction: {tokenizer.decode(outputs[0], skip_special_tokens=False)}")
+        
         return corrected_text
     except Exception as e:
         logging.error(f"Error in HTR Correction: {e}", exc_info=True)
@@ -34,12 +38,16 @@ def summarize_text(legal_text):
         inputs = tokenizer("summarize the following legal text: " + legal_text, return_tensors="pt", max_length=512, truncation=True)
         logging.debug(f"Tokenized Inputs for Summarization: {inputs}")
 
-        # Set do_sample=True
-        outputs = model.generate(**inputs, max_length=150, num_beams=4, early_stopping=True, temperature=0.8, do_sample=True)  
+        # Generate with beam search and sampling
+        outputs = model.generate(**inputs, max_length=150, num_beams=4, early_stopping=True, temperature=0.8, do_sample=True)
         logging.debug(f"Generated Summary (Tokens): {outputs}")
 
         summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
         logging.debug(f"Decoded Summary: {summary}")
+
+        # Re-tokenize the output for further inspection
+        logging.debug(f"Re-tokenized output for Summarization: {tokenizer.decode(outputs[0], skip_special_tokens=False)}")
+        
         return summary
     except Exception as e:
         logging.error(f"Error in Summarization: {e}", exc_info=True)
@@ -52,12 +60,16 @@ def answer_question(legal_text, question):
         inputs = tokenizer(formatted_input, return_tensors="pt", max_length=512, truncation=True)
         logging.debug(f"Tokenized Inputs for Question Answering: {inputs}")
 
-        # Set do_sample=True
-        outputs = model.generate(**inputs, max_length=150, num_beams=4, early_stopping=True, temperature=0.7, do_sample=True)  
+        # Generate with beam search and sampling
+        outputs = model.generate(**inputs, max_length=150, num_beams=4, early_stopping=True, temperature=0.7, do_sample=True)
         logging.debug(f"Generated Answer (Tokens): {outputs}")
 
         answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
         logging.debug(f"Decoded Answer: {answer}")
+
+        # Re-tokenize the output for further inspection
+        logging.debug(f"Re-tokenized output for Question Answering: {tokenizer.decode(outputs[0], skip_special_tokens=False)}")
+        
         return answer
     except Exception as e:
         logging.error(f"Error in Question Answering: {e}", exc_info=True)
@@ -68,6 +80,7 @@ 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.")
 
+    # Add the two clickable buttons with separate boxes and bold text
     with gr.Row():
         gr.HTML('''
             <div style="display: flex; gap: 10px;">
@@ -116,4 +129,4 @@ with gr.Blocks() as demo:
         clear_button.click(lambda: ("", "", ""), outputs=[legal_text_input_q, question_input, answer_output])
 
 # Launch the Gradio interface
-demo.launch()
+demo.launch()