import os
import gradio as gr
from transformers import T5Tokenizer, T5ForConditionalGeneration

# 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)

def correct_htr(raw_htr_text):
    # Tokenize the input text
    inputs = tokenizer(raw_htr_text, return_tensors="pt")

    # Generate corrected text
    outputs = model.generate(**inputs)
    corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return corrected_text

def summarize_text(legal_text):
    # Tokenize the input text with summarization prompt
    inputs = tokenizer("summarize: " + legal_text, return_tensors="pt")

    # Generate summary
    outputs = model.generate(**inputs)
    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return summary

def answer_question(legal_text, question):
    # Combine context and question
    inputs = tokenizer(f"question: {question} context: {legal_text}", return_tensors="pt")

    # Generate answer
    outputs = model.generate(**inputs)
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return answer

# 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.")

    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")
        correct_button.click(correct_htr, inputs=raw_htr_input, outputs=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")
        summarize_button.click(summarize_text, inputs=legal_text_input, outputs=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")
        answer_button.click(answer_question, inputs=[legal_text_input_q, question_input], outputs=answer_output)

demo.launch()