search

minimaxir / aitextgen

A robust Python tool for text-based AI training and generation using GPT-2.
https://docs.aitextgen.io
MIT License
1.84k stars 218 forks source link

Stop generating text after period #211

Open rmallick6806 opened 1 year ago

rmallick6806 commented 1 year ago

How do I get it to stop generating any more text after a marker like a period? Would be super helpful since it generates lots of incomplete sentences currently.

Vectorrent commented 1 year ago 
eos = ai.tokenizer.convert_tokens_to_ids(ai.tokenizer.tokenize('.')[0])
completion = ai.generate(
    eos_token_id=eos,
)

That will stop generation immediately upon hitting the eos (end-of-sentence) token.