import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# Load the local Tamil GPT-2 model
tokenizer = AutoTokenizer.from_pretrained("model")
model = AutoModelForCausalLM.from_pretrained("model")

generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

def generate_kavithai(prompt):
    if not prompt.strip():
        return "தயவுசெய்து ஒரு வரியை உள்ளிடவும்..."
    outputs = generator(prompt, max_length=50, num_return_sequences=1)
    return outputs[0]["generated_text"]

# Gradio UI setup
interface = gr.Interface(
    fn=generate_kavithai,
    inputs=gr.Textbox(lines=2, placeholder="உங்கள் எண்ணத்தை இங்கே எழுதுங்கள்...", label="📝 உங்கள் வரிகள்"),
    outputs=gr.Textbox(label="🎙️ கவிதை வெளியீடு"),
    title="தமிழ் கவிதை AI Bot ✍️",
    description="தமிழில் கவிதை உருவாக்கும் AI. உங்கள் வார்த்தைகளைப் பகிருங்கள் – ஒரு கவிதையை உருவாக்குவோம்!",
    theme="soft"
)

if __name__ == "__main__":
    interface.launch()