Update README.md

README.md

@@ -13,6 +13,52 @@ tags:
 - astronomy
 - astrophysics
 ---
-## Training procedure
 
-### Framework versions
+<p><h1>AstroLLaMA</h1></p>
+
+<p align="center">
+  <img src="https://huggingface.co/universeTBD/astrollama/resolve/main/images/astrollama-logo.png" alt="AstroLLaMA" width="500px"/>
+</p>
+
+## Loading the model
+
+```python
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained(
+    pretrained_model_name_or_path="universeTBD/astrollama"
+)
+model = AutoModelForCausalLM.from_pretrained(
+    pretrained_model_name_or_path="universeTBD/astrollama",
+    device_map="auto",
+)
+```
+
+## Generating text from a prompt
+
+```python
+import torch
+from transformers import pipeline
+
+generator = pipeline(
+    task="text-generation",
+    model=model,
+    tokenizer="tokenizer",
+    device_map="auto"
+)
+
+# Taken from https://arxiv.org/abs/2308.12823
+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. "
+
+# For reproducibility
+torch.manual_seed(42)
+
+generated_text = generator(
+    prompt,
+    do_sample=True,
+    max_length=512
+)
+```