import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load your fine-tuned model
model = AutoModelForSeq2SeqLM.from_pretrained("heyIamUmair/flan-t5-legal-finetuned_1st")
tokenizer = AutoTokenizer.from_pretrained("heyIamUmair/flan-t5-legal-finetuned_1st")

def generate_answer(question):
    prompt = f"Question: {question} \nAnswer:"
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)
    outputs = model.generate(**inputs, max_new_tokens=256)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Gradio interface
gr.Interface(
    fn=generate_answer,
    inputs=gr.Textbox(lines=2, placeholder="Ask your legal question here..."),
    outputs="text",
    title="📜 Pakistan Legal Assistant (FLAN-T5)",
    description="Ask any legal question related to Pakistani law. This assistant is fine-tuned on criminal, family, and property laws."
).launch()