Views
No views yet
1import torch
2from transformers import BloomTokenizerFast, BloomForCausalLM
3
4device = 'cuda' if torch.cuda.is_available() else 'cpu'
5ckpt = 'mrm8488/bloom-560m-finetuned-the-stack-rust'
6revision = '100k' # latest one at the moment
7
8tokenizer = BloomTokenizerFast.from_pretrained(ckpt)
9model = BloomForCausalLM.from_pretrained(ckpt, revision=revision).to(device)
10
11def complete_code(text):
12 inputs = tokenizer(text, return_tensors='pt')
13 input_ids = inputs.input_ids.to(device)
14 attention_mask = inputs.attention_mask.to(device)
15 output = model.generate(input_ids, attention_mask=attention_mask, max_length=2048, eos_token_id=tokenizer.eos_token_id)
16
17 return tokenizer.decode(output[0], skip_special_tokens=False)
18
19code_prompt = """
20use fastly::{Error, Request, Response};
21use serde_json::{json, Value};
22
23#[fastly::main]
24fn main(req: Request) -> Result<Response, Error> {
25 let mut response = req.send("origin_0")?;
26"""
27
28complete_code(code_prompt)@misc {manuel_romero_2022,
author = { {Manuel Romero} },
title = { bloom-560m-finetuned-the-stack-rust (Revision 5358462) },
year = 2022,
url = { https://huggingface.co/mrm8488/bloom-560m-finetuned-the-stack-rust },
doi = { 10.57967/hf/0236 },
publisher = { Hugging Face }
}