maahmad/lab_tinylm
1
0
Fork 0
generated from mwc/lab_tinylm
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Applying patch

Browse Source
This commit is contained in:
cplockport
2026-03-01 15:14:35 -05:00
parent c743298f05
commit 4840eafa5e
2 changed files with 15 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tlm/cli.py
View File

@@ -54,13 +54,13 @@ def generate(length, context_window_words, filepath, gutenberg, list_gutenberg,
raise click.UsageError("No training data provided. Must specify at least one of --filepath, --gutenberg, or --mbox.")
# Train and generate
model = TinyLanguageModel(n=context_window_words)
model.train(corpus)
tlm = TinyLanguageModel(n=context_window_words)
tlm.train(corpus)
if prompt:
prompt_tokens = tokenize_text(prompt, tokenize_opts)
else:
prompt_tokens = None
join_fn = ''.join if 'char' in tokenize_opts else None
join_fn = ''.join if 'char' in tokenize_opts else join
display_join = join_fn or join
if verbose:
@@ -70,18 +70,18 @@ def generate(length, context_window_words, filepath, gutenberg, list_gutenberg,
def step_callback(pattern, options, chosen):
opts = textwrap.fill(', '.join(sorted(set(options))), width=60)
rows.append([display_join(list(pattern)), opts, chosen])
output = model.generate(length, prompt=prompt_tokens, join_fn=join_fn, step_callback=step_callback)
output = tlm.generate(length, prompt=prompt_tokens, join_fn=join_fn, step_callback=step_callback)
click.echo(tabulate(rows, headers=["Context", "Options", "Selected"], tablefmt="simple"))
click.echo()
else:
output = model.generate(length, prompt=prompt_tokens, join_fn=join_fn)
output = tlm.generate(length, prompt=prompt_tokens, join_fn=join_fn)
click.echo(output)
if interact:
import code
code.interact(local=locals(), banner="Entering interactive shell. 'model' and 'output' are available.")
code.interact(local=locals(), banner="Entering interactive shell. 'tlm' and 'output' are available.")
if __name__ == "__main__":