Update README.md

README.md

@@ -32,22 +32,33 @@ AstroLLaMA-2-7B-Base_AIC is a specialized base language model for astronomy, dev
 
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
+# Load the model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("AstroMLab/astrollama-2-7b-base_aic")
 model = AutoModelForCausalLM.from_pretrained("AstroMLab/astrollama-2-7b-base_aic", device_map="auto")
 
-import torch
+# Create the pipeline with explicit truncation
 from transformers import pipeline
-
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
+generator = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device_map="auto",
+    truncation=True,
+    max_length=512
+)
 
 # Example prompt from an astronomy paper
 prompt = "In this letter, we report the discovery of the highest redshift, " \
     "heavily obscured, radio-loud QSO candidate selected using JWST NIRCam/MIRI, " \
     "mid-IR, sub-mm, and radio imaging in the COSMOS-Web field. "
 
-torch.manual_seed(42)  # For reproducibility
-generated_text = generator(prompt, do_sample=True, max_length=512)
+# Set seed for reproducibility
+torch.manual_seed(42)
+
+# Generate text
+generated_text = generator(prompt, do_sample=True)
 print(generated_text[0]['generated_text'])
 ```