skhatri/distilgpt2med

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')