Views
No views yet

1pip install -r requirements.txt
2pip install transformers datasets accelerate safetensors1from transformers import AutoTokenizer, AutoModelForCausalLM
2
3
4tokenizer = AutoTokenizer.from_pretrained("DireDreadlord/GemCod-codegen-270M")
5model = AutoModelForCausalLM.from_pretrained("DireDreadlord/GemCod-codegen-270M")
6model.to(device)
7model.eval()
8model.resize_token_embeddings(len(tokenizer))
9
10
11if tokenizer.pad_token is None:
12 tokenizer.pad_token = tokenizer.eos_token
13
14chat_template = """{% for message in messages %}{% if message['role'] == 'user' %}User: {{ message['content'] }}
15{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }}
16{% endif %}{% endfor %}"""
17tokenizer.chat_template = chat_template
18
19def generate_code(prompt, max_tokens) -> str:
20 messages = [
21 {
22 "role": "user",
23 "content": prompt
24 }
25 ]
26
27 formatted_prompt = tokenizer.apply_chat_template(
28 messages,
29 tokenize=False,
30 add_generation_prompt=True
31 )
32
33
34 inputs = tokenizer(formatted_prompt, return_tensors="pt").to(device)
35 input_length = inputs["input_ids"].shape[1]
36
37 with torch.no_grad():
38 outputs = model.generate(
39 **inputs,
40 max_new_tokens=256,
41 do_sample=False,
42 num_beams=1,
43 pad_token_id=tokenizer.eos_token_id,
44 eos_token_id=tokenizer.eos_token_id,
45 use_cache=False,
46 )
47
48 generated_tokens = outputs[0][input_length:]
49 generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
50 return generated_text
51
52
53prompt = "give me a cpp function that prints the first n fibonacci numbers"
54print("Prompt: ", prompt)
55
56result = generate_code(prompt)
57print(result)