ppisljar commited on
Commit
e2dd5f3
1 Parent(s): 3337766

Upload infer.py

Browse files
Files changed (1) hide show
  1. infer.py +30 -0
infer.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import onnxruntime
2
+ import torch
3
+
4
+ from transformers import AutoTokenizer
5
+
6
+ # setup GPU
7
+ if torch.cuda.is_available():
8
+ device = [0] # use 0th CUDA device
9
+ accelerator = 'gpu'
10
+ else:
11
+ device = 1
12
+ accelerator = 'cpu'
13
+
14
+ map_location = torch.device('cuda:{}'.format(device[0]) if accelerator == 'gpu' else 'cpu')
15
+
16
+ tokenizer = AutoTokenizer.from_pretrained('google/byt5-small')
17
+
18
+ sentence = "Kupil sem bicikel in mu zamenjal stol.".lower()
19
+
20
+ ort_session = onnxruntime.InferenceSession("g2p_t5.onnx", providers=["CPUExecutionProvider"])
21
+ input_ids = [sentence]
22
+ input_encoding = tokenizer(
23
+ input_ids, padding='longest', max_length=512, truncation=True, return_tensors='pt',
24
+ )
25
+ input_ids, attention_mask = input_encoding.input_ids, input_encoding.attention_mask
26
+ ort_inputs = {'input_ids': input_ids.numpy()}
27
+ ort_outs = ort_session.run(None, ort_inputs)
28
+ generated_ids = [ort_outs[0]]
29
+ generated_texts = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
30
+ print(generated_texts)