Update app.py (change to MT5Tokenizer; add button boxes and bold to button text)

app.py

@@ -1,40 +1,55 @@
 import os
 import gradio as gr
-from transformers import T5Tokenizer, T5ForConditionalGeneration
+from transformers import MT5Tokenizer, MT5ForConditionalGeneration
 
 # Load your fine-tuned mT5 model
 model_name = "Addaci/mT5-small-experiment-13-checkpoint-2790"
-tokenizer = T5Tokenizer.from_pretrained(model_name)
-model = T5ForConditionalGeneration.from_pretrained(model_name)
+tokenizer = MT5Tokenizer.from_pretrained(model_name)
+model = MT5ForConditionalGeneration.from_pretrained(model_name)
 
 def correct_htr(raw_htr_text):
-    inputs = tokenizer(raw_htr_text, return_tensors="pt", max_length=512, truncation=True)
-    outputs = model.generate(**inputs, max_length=128, num_beams=4, early_stopping=True)
-    corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return corrected_text
+    try:
+        inputs = tokenizer(raw_htr_text, return_tensors="pt", max_length=512, truncation=True)
+        outputs = model.generate(**inputs, max_length=128, num_beams=4, early_stopping=True)
+        corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return corrected_text
+    except Exception as e:
+        return str(e)
 
 def summarize_text(legal_text):
-    inputs = tokenizer("summarize: " + legal_text, return_tensors="pt", max_length=512, truncation=True)
-    outputs = model.generate(**inputs, max_length=150, num_beams=4, early_stopping=True)
-    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return summary
+    try:
+        inputs = tokenizer("summarize: " + legal_text, return_tensors="pt", max_length=512, truncation=True)
+        outputs = model.generate(**inputs, max_length=150, num_beams=4, early_stopping=True)
+        summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return summary
+    except Exception as e:
+        return str(e)
 
 def answer_question(legal_text, question):
-    formatted_input = f"question: {question} context: {legal_text}"
-    inputs = tokenizer(formatted_input, return_tensors="pt", max_length=512, truncation=True)
-    outputs = model.generate(**inputs, max_length=150, num_beams=4, early_stopping=True)
-    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return answer
+    try:
+        formatted_input = f"question: {question} context: {legal_text}"
+        inputs = tokenizer(formatted_input, return_tensors="pt", max_length=512, truncation=True)
+        outputs = model.generate(**inputs, max_length=150, num_beams=4, early_stopping=True)
+        answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return answer
+    except Exception as e:
+        return str(e)
 
-# Create the Gradio Blocks interface
+# Create the Gradio Blocks interface with boxed clickable buttons and bold text
 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.")
 
-    # Adding external link buttons at the top
+    # Adding external link buttons with a box around them and bold text
     with gr.Row():
-        gr.HTML('<a href="http://www.marinelives.org/wiki/Tools:_Admiralty_court_legal_glossary" target="_blank"><button>Admiralty Court Legal Glossary</button></a>')
-        gr.HTML('<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>HCA 13/70 Ground Truth (1654-55)</button></a>')
+        gr.HTML('''<div style="border: 2px solid black; padding: 10px; display: inline-block;">
+                      <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>
+                      <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>
+                   </div>''')
 
     with gr.Tab("Correct HTR"):
         gr.Markdown("### Correct Raw HTR Text")
@@ -70,4 +85,6 @@ with gr.Blocks() as demo:
 # Launch the Gradio interface
 demo.launch()
 
+   
+