| from __future__ import annotations |
|
|
| import argparse |
| import os |
| from pathlib import Path |
|
|
| import gradio as gr |
| from tokenizers import Tokenizer |
|
|
| from nayara.generate import format_prompt, load_model, stream_generate |
|
|
|
|
| def main() -> None: |
| p = argparse.ArgumentParser(description="NAYARA local code assistant") |
| p.add_argument("checkpoint") |
| p.add_argument("--host", default="127.0.0.1") |
| p.add_argument("--port", type=int, default=7860) |
| p.add_argument("--threads", type=int, default=max(1, (os.cpu_count() or 4) // 2)) |
| args = p.parse_args() |
| root = Path(args.checkpoint) |
| tokenizer = Tokenizer.from_file(str(root / "tokenizer.json")) |
| model = load_model(root, int8=True, threads=args.threads) |
|
|
| def answer(prompt: str, max_tokens: int, temperature: float): |
| for partial in stream_generate(model, tokenizer, format_prompt(prompt), int(max_tokens), temperature): |
| yield partial |
|
|
| with gr.Blocks(title="NAYARA v1.0", theme=gr.themes.Soft()) as demo: |
| gr.Markdown("# NAYARA v1.0\nLocal, from-scratch code intelligence.") |
| prompt = gr.Textbox(label="Ask or start coding", lines=10, placeholder="Write a fast Python LRU cache...") |
| with gr.Row(): |
| max_tokens = gr.Slider(16, 512, value=192, step=16, label="Max new tokens") |
| temperature = gr.Slider(0, 1.2, value=0.2, step=0.05, label="Temperature") |
| run = gr.Button("Generate", variant="primary") |
| output = gr.Code(label="NAYARA", language=None, lines=16) |
| run.click(answer, [prompt, max_tokens, temperature], output) |
| prompt.submit(answer, [prompt, max_tokens, temperature], output) |
| demo.queue(default_concurrency_limit=1).launch(server_name=args.host, server_port=args.port) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
|
|
