Update README.md

README.md

@@ -3,6 +3,8 @@
 
 This model is a pretrained autoregressive transformer model in GPT-style, trained on a large number of protein sequences.  
 
+Dataset: https://huggingface.co/datasets/lamm-mit/GPTProteinPretrained
+
 Load pretrained model:
 
 ```python
@@ -26,7 +28,9 @@ model.config.use_cache = False
 Sample inference using the "Sequence<...>" task, where here, the model will simply autocomplete the sequence starting with "AIIAA":
 
 ```python
-prompt = "Sequence<AIIAA"
+import torch 
+device='cuda'
+prompt = "Sequence<ETAVPKLLQAL"
 generated = torch.tensor(tokenizer.encode(prompt, add_special_tokens = False)) .unsqueeze(0).to(device)
 print(generated.shape, generated)
 
@@ -35,9 +39,9 @@ sample_outputs = model.generate(
                                 eos_token_id =tokenizer.eos_token_id,
                                 do_sample=True,   
                                 top_k=500, 
-                                max_length = 300,
+                                max_length = 1024,
                                 top_p=0.9, 
-                                num_return_sequences=3,
+                                num_return_sequences=6,
                                 temperature=1,
                                 ).to(device)