Update README.md

README.md

@@ -96,7 +96,21 @@ We provide step-by-step reproducation tutorials for tech enthusiasts, AI practit
 ## LLM360 Developer Suite
 We provide step-by-step finetuning tutorials for tech enthusiasts, AI practitioners and academic or industry researchers here [llm360.ai/pretraining].
 
-# Put inference section here.
+# Loading K2
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("LLM360/K2")
+model = AutoModelForCausalLM.from_pretrained("LLM360/K2")
+
+prompt = 'int add(int x, int y) {'
+
+input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+gen_tokens = model.generate(input_ids, do_sample=True, max_length=128)
+
+print("-"*20 + "Output for model"  + 20 * '-')
+print(tokenizer.batch_decode(gen_tokens)[0])
+```