Total rewrite of app.py (Moved to Inference Client approach)

app.py

@@ -1,78 +1,47 @@
-import os
+# Cell 1B: Inference Client
+
 import gradio as gr
+from huggingface_hub import InferenceClient
 import logging
-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 = MT5Tokenizer.from_pretrained(model_name)
-model = MT5ForConditionalGeneration.from_pretrained(model_name)
+# Initialize Inference Client
+client = InferenceClient(model="Addaci/mT5-small-experiment-13-checkpoint-2790")
 
 def correct_htr(raw_htr_text):
     try:
-        logging.info("Processing HTR correction...")
-        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}")
-
-        # 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
+        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
     except Exception as e:
-        logging.error(f"Error in HTR Correction: {e}", exc_info=True)
+        logging.error(f"Error in HTR correction: {e}", exc_info=True)
         return str(e)
 
 def summarize_text(legal_text):
     try:
-        logging.info("Processing summarization...")
-        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}")
-
-        # 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
+        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
     except Exception as e:
-        logging.error(f"Error in Summarization: {e}", exc_info=True)
+        logging.error(f"Error in summarization: {e}", exc_info=True)
         return str(e)
 
 def answer_question(legal_text, question):
     try:
-        logging.info("Processing question-answering...")
+        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}"
-        inputs = tokenizer(formatted_input, return_tensors="pt", max_length=512, truncation=True)
-        logging.debug(f"Tokenized Inputs for Question Answering: {inputs}")
-
-        # 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
+        result = client.text_generation(formatted_input)
+        logging.debug(f"Generated answer: {result}")
+        return result['generated_text']  # Extracting the generated text from the response
     except Exception as e:
-        logging.error(f"Error in Question Answering: {e}", exc_info=True)
+        logging.error(f"Error in question-answering: {e}", exc_info=True)
         return str(e)
 
 # Create the Gradio Blocks interface
@@ -80,7 +49,6 @@ 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;">