Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2import torch
3
4prompt = '''def min(arr):
5 """
6 Returns the minimum value from the list `arr`.
7
8 Parameters:
9 - arr (list): A list of numerical values.
10
11 Returns:
12 - The minimum value in `arr`.
13 """
14'''
15
16tokenizer = AutoTokenizer.from_pretrained(
17 "TechxGenus/Mini-Jamba",
18 trust_remote_code=True,
19)
20tokenizer.pad_token = tokenizer.eos_token
21model = AutoModelForCausalLM.from_pretrained(
22 "TechxGenus/Mini-Jamba",
23 torch_dtype=torch.float16,
24 device_map="auto",
25 trust_remote_code=True,
26)
27inputs = tokenizer.encode(prompt, return_tensors="pt")
28outputs = model.generate(
29 input_ids=inputs.to(model.device),
30 max_new_tokens=64,
31 do_sample=False,
32)
33print(tokenizer.decode(outputs[0], skip_special_tokens=True))
34