Views
No views yet

1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3model_path="OpenCodeInterpreter-DS-6.7B"
4
5tokenizer = AutoTokenizer.from_pretrained(model_path)
6model = AutoModelForCausalLM.from_pretrained(
7 model_path,
8 torch_dtype=torch.bfloat16,
9 device_map="auto",
10)
11model.eval()
12
13prompt = "Write a function to find the shared elements from the given two lists."
14inputs = tokenizer.apply_chat_template(
15 [{'role': 'user', 'content': prompt }],
16 return_tensors="pt"
17 ).to(model.device)
18outputs = model.generate(
19 inputs,
20 max_new_tokens=1024,
21 do_sample=False,
22 pad_token_id=tokenizer.eos_token_id,
23 eos_token_id=tokenizer.eos_token_id,
24)
25print(tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True))