Update README.md

README.md

@@ -46,3 +46,58 @@ python finetune.py \
 --eval-file code_eval.jsonl --wandb-project jerboa --wandb-log-model \
 --wandb-watch gradients --num-epochs 2
 ```
+
+
+```Python
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+
+TOKENIZER_SOURCE = 'tiiuae/falcon-40b'
+BASE_MODEL = 'jinaai/falcon-40b-code-alpaca'
+DEVICE = "cuda"
+
+PROMPT = """
+Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
+
+### Instruction:
+Write a for loop in python
+
+### Input:
+
+### Response:
+"""
+model = AutoModelForCausalLM.from_pretrained(
+    pretrained_model_name_or_path=BASE_MODEL,
+    torch_dtype=torch.float16,
+    trust_remote_code=True,
+    device_map='auto',
+)
+
+model.eval()
+
+tokenizer = AutoTokenizer.from_pretrained(
+    TOKENIZER_SOURCE,
+    trust_remote_code=True,
+    padding_side='left',
+)
+tokenizer.pad_token = tokenizer.eos_token
+
+inputs = tokenizer(PROMPT, return_tensors="pt")
+input_ids = inputs["input_ids"].to(DEVICE)
+input_attention_mask = inputs["attention_mask"].to(DEVICE)
+
+with torch.no_grad():
+    generation_output = model.generate(
+        input_ids=input_ids,
+        attention_mask=input_attention_mask,
+        return_dict_in_generate=True,
+        max_new_tokens=32,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+generation_output = generation_output.sequences[0]
+output = tokenizer.decode(generation_output, skip_special_tokens=True)
+
+print(output)
+
+```