Skip to content

EOS String Implementation Summary

This document summarizes the implementation of the custom eos_string parameter feature.

Changes Made

1. Core Generation Module (steadytext/core/generator.py)

  • Added eos_string parameter to DeterministicGenerator.generate() method
  • Default value: "[EOS]" (special marker for model's default EOS token)
  • When custom value provided, it's added to the stop sequences

  • Added eos_string parameter to DeterministicGenerator.generate_iter() method

  • Supports streaming generation with custom stop strings
  • Added include_logprobs parameter for compatibility with CLI

  • Updated caching logic to include eos_string in cache key when not default

  • Ensures different eos_strings produce separately cached results

2. Public API (steadytext/__init__.py)

  • Updated generate() function signature:

    def generate(prompt: str, return_logprobs: bool = False, eos_string: str = "[EOS]")
    

  • Updated generate_iter() function signature:

    def generate_iter(prompt: str, eos_string: str = "[EOS]", include_logprobs: bool = False)
    

3. CLI Updates

Generate Command (steadytext/cli/commands/generate.py)

  • Added --eos-string parameter (default: "[EOS]")
  • Passes eos_string to both batch and streaming generation

Main CLI (steadytext/cli/main.py)

  • Added --quiet / -q flag to silence log output
  • Sets logging level to ERROR for both steadytext and llama_cpp loggers when quiet mode is enabled

4. Tests (tests/test_steadytext.py)

Added three new test methods: - test_generate_with_custom_eos_string() - Tests basic eos_string functionality - test_generate_iter_with_eos_string() - Tests streaming with custom eos_string - test_generate_eos_string_with_logprobs() - Tests combination of eos_string and logprobs

5. Test Scripts

Created two test scripts for manual verification: - test_eos_string.py - Python script testing various eos_string scenarios - test_cli_eos.sh - Bash script testing CLI functionality

Usage Examples

Python API

import steadytext

# Use model's default EOS token
text = steadytext.generate("Hello world", eos_string="[EOS]")

# Stop at custom string
text = steadytext.generate("List items until END", eos_string="END")

# Streaming with custom eos
for token in steadytext.generate_iter("Generate text", eos_string="STOP"):
    print(token, end="")

CLI

# Default behavior
steadytext "Generate some text"

# Custom eos string
steadytext "Generate until DONE" --eos-string "DONE"

# Quiet mode (no logs)
steadytext --quiet "Generate without logs"

# Streaming with custom eos
steadytext "Stream until END" --stream --eos-string "END"

Implementation Notes

  1. The "[EOS]" string is a special marker that tells the system to use the model's default EOS token and stop sequences.

  2. When a custom eos_string is provided, it's added to the existing stop sequences rather than replacing them.

  3. Cache keys include the eos_string when it's not the default, ensuring proper caching behavior.

  4. The quiet flag affects all loggers in the steadytext namespace and llama_cpp if present.