import gradio as gr
import torch
import re
from transformers import BartForConditionalGeneration, PreTrainedTokenizerFast

# Load Model
MODEL_PATH = "./model"
TOKENIZER_PATH = "./tokenizer"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = PreTrainedTokenizerFast.from_pretrained(TOKENIZER_PATH)
model = BartForConditionalGeneration.from_pretrained(MODEL_PATH).to(device)
model.eval()

def paraphrase(text):
    if not text.strip(): return ""
    # Basic validation
    if re.search(r'[a-zA-Z]', text): return "Error: Please provide input only in Tamil."
    
    inputs = tokenizer(text, return_tensors="pt", max_length=128, truncation=True).to(device)
    with torch.no_grad():
        outputs = model.generate(inputs["input_ids"], max_new_tokens=60, num_beams=5, early_stopping=True)
    
    result = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
    return re.sub(r'[^\u0B80-\u0BFF\s.,!?-]', '', result).strip() + "."

# UI Setup
demo = gr.Interface(
    fn=paraphrase,
    inputs=gr.Textbox(label="Input (உள்ளீடு)"),
    outputs=gr.Textbox(label="Paraphrase (மாற்று வாக்கியம்)"),
    title="Tamil Paraphrase AI",
    article="Disclaimer: This is an AI and can make mistakes."
)

demo.launch()