Upload 2 files

main.py

@@ -0,0 +1,33 @@
+from transformers import GPT2Tokenizer
+import torch
+import pickle
+
+tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+tokenizer.pad_token = tokenizer.eos_token
+#model = GPT2LMHeadModel.from_pretrained('gpt2')
+#model.load_state_dict(torch.load("/kaggle/input/poem-pt/poem_model.pt"))
+#torch.save(model,"poem_model.pt")
+model = torch.load("poem_model.pt")
+
+def infer(inp):
+    inp = tokenizer(inp,return_tensors="pt")
+    X = inp["input_ids"] #.to(device)
+    a = inp["attention_mask"] #.to(device)
+    output = model.generate(X, 
+                            attention_mask=a,
+                            max_length=100,
+                            min_length=10,
+                            early_stopping=True,
+                            num_beams=5, 
+                            no_repeat_ngram_size=2)
+    
+    output = tokenizer.decode(output[0])
+    
+    return output
+
+#pickle.dump(model,open('poem_model.pt','wb'))
+
+output = infer(" I shall go")
+
+
+print(output)

poem_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e0d6b10f3fe4fba2c6a28894281e49ed41991ca371dcb26ae84b4485119561e1
+size 510424627