feat: improve example

README.md

@@ -10,6 +10,32 @@ Program synthesis strives to generate a computer program as a solution to a give
 
 # How to use
 ```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
 tokenizer = AutoTokenizer.from_pretrained("shpotes/codegen-350M-mono")
 model = AutoModelForCausalLM.from_pretrained("shpotes/codegen-350M-mono", trust_remote_code=True)
+
+input_ids = tokenizer(
+    context,
+    truncation=True,
+    padding=True,
+    return_tensors='pt',
+    pad_token_id=pad_token_id,
+).input_ids
+
+input_ids_len = input_ids.shape[1]
+
+with torch.no_grad():
+  input_ids = input_ids
+  tokens = model.generate(
+      input_ids,
+      do_sample=True,
+      num_return_sequences=num_return_sequences,
+      temperature=temp,
+      max_length=input_ids_len + max_length_sample,
+      top_p=top_p,
+      use_cache=True,
+  )
+  text = tokenizer.batch_decode(tokens[:, input_ids_len:, ...])
 ```