import gradio as gr from transformers import AutoModelForSeq2SeqLM, AutoTokenizer import logging # Setup logging (optional, but helpful for debugging) logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s') # Load the Flan-T5 Small model and tokenizer model_id = "google/flan-t5-small" tokenizer = AutoTokenizer.from_pretrained(model_id) 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...") 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) # Cap max_length outputs = model.generate(**inputs, max_length=max_length, 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." 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...") 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) # Cap max_length outputs = model.generate(**inputs, max_length=max_length, 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." 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...") 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) # Cap max_length outputs = model.generate(**inputs, max_length=max_length, 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 "", "", "" # Create the Gradio Blocks interface 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('''
''') # Sliders for max_new_tokens and temperature with gr.Row(): max_new_tokens_slider = gr.Slider(minimum=10, maximum=512, value=128, step=1, label="Max New Tokens") temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature") 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_button.click(correct_htr, inputs=[raw_htr_input, max_new_tokens_slider, temperature_slider], outputs=corrected_output) clear_button.click(clear_fields, outputs=[raw_htr_input, corrected_output]) 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_button.click(summarize_text, inputs=[legal_text_input, max_new_tokens_slider, temperature_slider], outputs=summary_output) clear_button.click(clear_fields, outputs=[legal_text_input, summary_output]) 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...") question_input = gr.Textbox(lines=2, placeholder="Enter your question...") 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, max_new_tokens_slider, temperature_slider], outputs=answer_output) clear_button.click(clear_fields, outputs=[legal_text_input_q, question_input, answer_output]) # Model warm-up (optional, but useful for performance) model.generate(**tokenizer("Warm-up", return_tensors="pt"), max_length=10) # Launch the Gradio interface if __name__ == "__main__": demo.launch()