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() |