Create README.md

README.md

@@ -0,0 +1,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')
+```