File size: 748 Bytes
cff32ca
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
Code to test this model.

```
import torch
import time

device_name="cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device_name)

model_name="skhatri/distilgpt2med"
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
model.to(device)

raw_input = "Headache Cough"

import sys
if len(sys.argv) > 1:
  raw_input = sys.argv[1]

start=time.time()

input_ids = tokenizer.encode(raw_input, return_tensors='pt').to(device)
output = model.generate(input_ids)

response = tokenizer.decode(output[0], skip_special_tokens=True)
print(response)
end=time.time()
print(f'Time taken: {round(end - start, 2)} seconds')
```