Views
No views yet
1from transformers import AutoTokenizer, GemmaForCausalLM
2import torch
3import time
4
5def inference(input_text):
6 start_time = time.time()
7 input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)
8 input_length = input_ids["input_ids"].shape[1]
9 outputs = model.generate(
10 input_ids=input_ids["input_ids"],
11 max_length=1024,
12 do_sample=False)
13 generated_sequence = outputs[:, input_length:].tolist()
14 res = tokenizer.decode(generated_sequence[0])
15 end_time = time.time()
16 return {"output": res, "latency": end_time - start_time}
17
18model_id = "NexaAIDev/Octopus-v2"
19tokenizer = AutoTokenizer.from_pretrained(model_id)
20model = GemmaForCausalLM.from_pretrained(
21 model_id, torch_dtype=torch.bfloat16, device_map="auto"
22)
23
24input_text = "Take a selfie for me with front camera"
25nexa_query = f"Below is the query from the users, please call the correct function and generate the parameters to call the function.\n\nQuery: {input_text} \n\nResponse:"
26start_time = time.time()
27print("nexa model result:\n", inference(nexa_query))
28print("latency:", time.time() - start_time," s")