Update README

README.md

@@ -30,22 +30,25 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Load the tokenizer and model
-model_path = "nvidia/Minitron-4B-Base"
-tokenizer = AutoTokenizer.from_pretrained(model_path)
+model_path = 'nvidia/Minitron-4B-Base'
+tokenizer  = AutoTokenizer.from_pretrained(model_path)
 
-device='cuda'
-dtype=torch.bfloat16
-model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype, device_map=device)
+device = 'cuda'
+dtype  = torch.bfloat16
+model  = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype, device_map=device)
 
 # Prepare the input text
-prompt = "To be or not to be,"
-input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
+prompt = 'Complete the paragraph: our solar system is'
+inputs = tokenizer.encode(prompt, return_tensors='pt').to(model.device)
 
 # Generate the output
-output_ids = model.generate(input_ids, max_length=50, num_return_sequences=1)
+outputs = model.generate(inputs,
+                         max_length=20,
+                         num_return_sequences=1,
+                         pad_token_id=tokenizer.eos_token_id)
 
 # Decode and print the output
-output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
 print(output_text)
 ```