Update README.md

README.md

@@ -58,11 +58,15 @@ GPT-J learns an inner representation of the English language that can be used to
 
 ```python
 # libraries and a wrapper around hivemind's quantization code
-!pip install transformers==4.14.1 bitsandbytes-cuda111==0.26.0 git+https://github.com/aicrumb/gpt-j-8bit -q
-import gptj_8bit as gptj
+!pip install transformers==4.14.1 bitsandbytes-cuda111==0.26.0 git+https://github.com/aicrumb/transformers-8bit -q
+import transformers_8bit
 
-model, tokenizer, config = gptj.load("crumb/gpt-j-6b-shakespeare", device='cuda')
-gptj.generate(model, tokenizer, "Romeo:", min_length=64, max_length=64)
+model, tokenizer, config = transformers_8bit.load_gptj("crumb/gpt-j-6b-shakespeare", device='cuda')
+
+prompt = tokenizer("Romeo:", return_tensors='pt')
+prompt = {key: value.to('cuda') for key, value in prompt.items()}
+out = model.generate(**prompt, min_length=64, max_length=64, do_sample=True, pad_token_id=tokenizer.eos_token_id)
+print(tokenizer.decode(out[0]))
 
 """ example output
 Romeo: [Aside] And but in night, how tedious