File size: 1,658 Bytes
a58ece3 | 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 | """CLI: interactive Go code generation."""
import click
import tensorflow as tf
from llm_go.model.transformer import GoLLM
from llm_go.tokenizer.go_tokenizer import GoTokenizer
@click.command()
@click.option("--model-dir", required=True, help="Checkpoint directory")
@click.option("--tok-dir", required=True, help="Tokenizer directory")
@click.option("--prompt", default="", help="Prompt text (leave empty for interactive)")
@click.option("--max-tokens", default=256, show_default=True)
@click.option("--temperature", default=0.8, show_default=True)
@click.option("--top-p", default=0.95, show_default=True)
@click.option("--top-k", default=50, show_default=True)
def main(model_dir, tok_dir, prompt, max_tokens, temperature, top_p, top_k):
"""Generate Go code with a trained GoLLM checkpoint."""
tok = GoTokenizer.load(tok_dir)
model = GoLLM.from_pretrained(model_dir)
def generate(text: str) -> str:
ids = tok.encode(text)
out = model.generate(
tf.constant([ids], dtype=tf.int32),
max_new_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
top_k=top_k,
)
return tok.decode(out[0].numpy().tolist())
if prompt:
click.echo(generate(prompt))
else:
click.echo("GoLLM interactive shell. Type 'exit' to quit.\n")
while True:
try:
text = click.prompt(">>> ")
except (EOFError, KeyboardInterrupt):
break
if text.strip().lower() in {"exit", "quit"}:
break
click.echo(generate(text))
|