import gradio as gr
from train import train
from test_model import generate_response

# Create the training interface
training_interface = gr.Interface(
    fn=train,
    inputs=None,
    outputs="text",
    title="Llama 2 AA Training",
    description="Fine-tune Llama 2 on AA dataset using LoRA"
)

# Create the testing interface
testing_interface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(label="Ask a question about AA"),
    outputs=gr.Textbox(label="Response"),
    title="Test AA Fine-tuned Llama 2",
    description="Ask questions about AA to test the fine-tuned model"
)

# Combine both interfaces
demo = gr.TabbedInterface(
    [training_interface, testing_interface],
    tab_names=["Train", "Test"]
)

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