aa_book / app.py
lambchop11's picture
Fix import error and rename test file
419b514
raw
history blame contribute delete
780 Bytes
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()