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

# Load model and tokenizer from Hugging Face model repository
model_path = "EzioDevio/iraqi_dialect_llm"  # Use the Hugging Face Hub model path

# Load the model and tokenizer directly from Hugging Face
model = AutoModelForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Define the text generation function
def generate_text(input_text):
    generated = text_generator(input_text, max_length=50, num_return_sequences=1)
    return generated[0]["generated_text"]

# Set up the Gradio interface
iface = gr.Interface(
    fn=generate_text,
    inputs="text",
    outputs="text",
    title="Iraqi Dialect Language Model",
    description="Generate text in Iraqi Arabic dialect.",
    examples=["شلونك اليوم؟", "وين رايح؟", "صباح الخير"]
)

# Launch the app
if __name__ == "__main__":
    iface.launch()