Update README.md

README.md

@@ -16,6 +16,39 @@ language:
 
 The Stack contains over 6TB of permissively-licensed source code files covering 358 programming languages. The dataset was created as part of the BigCode Project, an open scientific collaboration working on the responsible development of Large Language Models for Code (Code LLMs). The Stack serves as a pre-training dataset for Code LLMs, i.e., code-generating AI systems which enable the synthesis of programs from natural language descriptions as well as other from code snippets.
 
+## Example of usage
+
+```py
+import torch
+from transformers import BloomTokenizerFast, BloomForCausalLM
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+ckpt = 'mrm8488/bloom-560m-finetuned-the-stack-rust'
+revision = '100k' #latest one at the moment 
+
+tokenizer = BloomTokenizerFast.from_pretrained(ckpt)
+model = BloomForCausalLM.from_pretrained(ckpt, revision=revision).to(device)
+
+def complete_code(text):
+    inputs = tokenizer(text, return_tensors='pt')
+    input_ids = inputs.input_ids.to(device)
+    attention_mask = inputs.attention_mask.to(device)
+    output = model.generate(input_ids, attention_mask=attention_mask, max_length=2048, eos_token_id=tokenizer.eos_token_id)
+
+    return tokenizer.decode(output[0], skip_special_tokens=False)
+    
+code_prompt = """
+use fastly::{Error, Request, Response};
+use serde_json::{json, Value};
+
+#[fastly::main]
+fn main(req: Request) -> Result<Response, Error> {
+  let mut response = req.send("origin_0")?;
+"""
+
+complete_code(code_prompt)
+```
+