stefstefstef/commit-message

1import gradio as gr2from transformers import AutoModelForCausalLM, AutoTokenizer3 4gr.load("models/bigcode/starcoder2-7b").launch()5 6checkpoint = "bigcode/starcoder2-7b"7device = "cuda" # for GPU usage or "cpu" for CPU usage8 9tokenizer = AutoTokenizer.from_pretrained(checkpoint)10# for multiple GPUs install accelerate and do `model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto")`11model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)12 13inputs = tokenizer.encode("def print_hello_world():", return_tensors="pt").to(device)14outputs = model.generate(inputs)15print(tokenizer.decode(outputs[0]))