Making this library more like Hugging Face

[oobabooga]

I have expressed my interest in having RWKV officially implemented in Hugging Face in huggingface/transformers#17230.

Meanwhile, I have a distilled set of suggestions for how this library could be made more familiar to people who are already used to transformers and AutoModelForCausalLM.

Maybe some of these are already possible in the current version of rwkv. If so, I would be grateful if you could let me know how.

1. Being able to load the tokenizer explicitly

with something like

tokenizer = RWKVTokenizer.from_pretrained("/path/to/20B_tokenizer.json")

and then use it with

prompt = "Hello, my name is "
input_ids = tokenizer.encode(prompt)

Having the ability to count the number of tokens in a given prompt is very useful.

2. Generating text with input_ids as input rather than a string

Something like

output_ids = model.generate(input_ids, temperature=0.8, top_p=0.95)
output_text = tokenizer.decode(output_ids)

3. Generation parameters

Many parameters are available for model.generate() in HF, but it seems to me that the absolutely essential ones that everyone uses are:

1. temperature ✅
2. top_p ✅
3. top_k
4. repetition_penalty

I am aware that alpha_frequency and alpha_presence are implemented, but these parameters are not usually found in presets that people have already come up with while working with other models. For this reason, having repetition_penalty would be valuable.

[oobabooga]

For reference, I am currently using this wrapper to run the model:

https://github.com/oobabooga/text-generation-webui/blob/main/modules/RWKV.py

[BlinkDL]

from tokenizers import Tokenizer # pip install tokenizers
tokenizer = Tokenizer.from_file("20B_tokenizer.json")
print(tokenizer.encode("Hello World").ids)
print(tokenizer.decode([12092, 3645])

output_ids = tokenizer.encode(model.generate(tokenizer.decode(input_ids), temperature=0.8, top_p=0.95)).ids
output_text = tokenizer.decode(output_ids)

3. soon :)

[oobabooga]

Thanks, based on your suggestions I have added a RWKVTokenizer wrapper: oobabooga/text-generation-webui@e91f4bc