Update README.md

README.md

@@ -4,6 +4,9 @@ datasets:
 - flytech/python-codes-25k
 tags:
 - code
+language:
+- en
+library_name: transformers
 ---
 # GPT2 PyCode
 
@@ -51,11 +54,15 @@ Use the code below to get started with the model.
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import re
 
-tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
-model = GPT2Model.from_pretrained('gpt2')
 
-text = "Replace me by any text you'd like."
-encoded_input = tokenizer.encode_plus(f"<sos><user>{prompt}</user><assistant>", max_length=20, truncation=True, return_tensors="pt")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+tokenizer = GPT2Tokenizer.from_pretrained('maharnab/gpt2_pycode')
+model = GPT2LMHeadModel.from_pretrained('maharnab/gpt2_pycode')
+model.to(device)
+
+prompt = "How to reverse a string in Python."
+encoded_input = tokenizer.encode_plus(f"<sos><user>{prompt}</user><assistant>", max_length=20, truncation=True, return_tensors="pt").to(device)
 
 input_ids = encoded_input['input_ids']
 attention_mask = encoded_input['attention_mask']