Update README.md

README.md

@@ -37,7 +37,37 @@ AstroLLaMA-3-8B is a specialized base language model for astronomy, developed by
 
 ## Generating text from a prompt
 
-[The code example remains the same as in the previous version]
+```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")
+
+# Create the pipeline with explicit truncation
+from transformers import pipeline
+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. "
+
+# Set seed for reproducibility
+torch.manual_seed(42)
+
+# Generate text
+generated_text = generator(prompt, do_sample=True)
+print(generated_text[0]['generated_text'])
+```
 
 ## Model Limitations and Biases