Spaces:
Sleeping
Sleeping
File size: 2,578 Bytes
f907cd1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | import gradio as gr
from aber_llm.config import AberConfig
from aber_llm.service import AberLLMService
config = AberConfig()
service = AberLLMService(config=config)
def generate_text(prompt, max_new_tokens, temperature, top_k):
return service.generate(
prompt=prompt,
max_new_tokens=int(max_new_tokens),
temperature=float(temperature),
top_k=int(top_k),
)
def train_model(extra_text, steps):
return service.train(extra_text=extra_text, steps=int(steps))
def reset_model():
return service.reset()
with gr.Blocks(
title="aber Small Model",
theme=gr.themes.Soft(primary_hue="green", secondary_hue="blue"),
) as demo:
gr.Markdown(
"""
# aber
An improved small language model written in Python from scratch.
- Model name: `aber`
- No external pretrained LLM
- Word-level tokenizer
- GRU language model
- Local CPU training and generation
"""
)
with gr.Tab("Generate"):
prompt_input = gr.Textbox(
label="Prompt",
value="User: hello\naber:",
lines=6,
)
with gr.Row():
max_tokens_input = gr.Slider(10, 160, value=72, step=2, label="Max New Tokens")
temperature_input = gr.Slider(0.2, 1.3, value=0.75, step=0.05, label="Temperature")
top_k_input = gr.Slider(1, 20, value=8, step=1, label="Top-K")
generate_button = gr.Button("Generate", variant="primary")
output_text = gr.Textbox(label="Output", lines=10)
output_status = gr.Textbox(label="Status", lines=4)
with gr.Tab("Train"):
extra_text_input = gr.Textbox(
label="Extra Training Text",
placeholder="Add more local text to train aber on your own data.",
lines=10,
)
steps_input = gr.Slider(10, 400, value=120, step=10, label="Training Steps")
train_button = gr.Button("Train / Continue Training", variant="primary")
reset_button = gr.Button("Reset aber")
train_status = gr.Textbox(label="Training Status", lines=6)
generate_button.click(
fn=generate_text,
inputs=[prompt_input, max_tokens_input, temperature_input, top_k_input],
outputs=[output_text, output_status],
)
train_button.click(
fn=train_model,
inputs=[extra_text_input, steps_input],
outputs=[train_status],
)
reset_button.click(
fn=reset_model,
outputs=[train_status],
)
if __name__ == "__main__":
demo.launch()
|