import gradio as gr
from app.model.model import NigerianLanguageModel
from app.model.config import ModelConfig

# Initialize model
config = ModelConfig()
model = NigerianLanguageModel(config)

# Create interface
def generate_text(prompt: str, language: str) -> str:
    tagged_prompt = f"[{language.upper()}] {prompt}"
    return model.generate(tagged_prompt)

# Define Gradio interface
interface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(label="Enter your prompt"),
        gr.Dropdown(choices=["YORUBA", "IGBO", "HAUSA"], label="Select Language")
    ],
    outputs=gr.Textbox(label="Generated Text"),
    title="Nigerian Language Generator",
    description="Generate text in Yoruba, Igbo, or Hausa using a fine-tuned GPT model."
)

# Start the interface
if __name__ == "__main__":
    interface.launch()