Astrophysics Prompting with OLMo

Astrophysics Prompting with OLMo#

At University of Washington, eScience Institute, we are very fortunate to have an in-house astronomer, Professor Andy Connolly.

He has been generously sharing his knowledge and came up with a very interesting set of questions for us to try out with OLMo:

  • What is dark matter?

  • How many dimensions are there in the elemental abundances of stars?

  • What is the evidence for cosmological parity violation?

  • What observations show evidence for dark matter?

  • What are the most popular theories for what dark matter might be made of?

  • Does dark matter need to be a particle?

  • What is the expected mass of dark matter?

  • What experiments are currently ongoing and searching for dark matter?

  • How do the above experiments work?

  • What type of dark matter would the LHC be able to detect?

  • Can you modify gravity to explain dark matter? If so how?

  • How do astronomers detect the presence of planets around nearby stars?

  • What is an eclipsing binary star?

  • Can you describe a taxonomy for different classes of variable stars?

  • What would the light curve of a tidal disruption event look like? How would it be different from a supernova?

  • How do supernova explode?

  • How is a Type II SN different from a Type 1a?

Let’s analyze how OLMo does without any additional context with these questions or your own specific domain questions.

Set up the OLMo Model and Prompt#

We’ll begin with a recap of the previous module, setting up the OLMo model and prompt.

from langchain_community.llms import LlamaCpp
from langchain_core.callbacks import StreamingStdOutCallbackHandler
from ssec_tutorials import download_olmo_model

OLMO_MODEL = download_olmo_model()

Model already exists at /Users/lsetiawan/.cache/ssec_tutorials/OLMo-7B-Instruct-Q4_K_M.gguf

This time during the model setup, we’ll try to increase the n_ctx, input context length, to 2048 tokens and the max_tokens, maximum tokens generated by the model, to 512 tokens. This is so later we can really expand on the questions that we ask the model and get a more expansive answer.

olmo = LlamaCpp(
    model_path=str(OLMO_MODEL),
    callbacks=[StreamingStdOutCallbackHandler()],
    temperature=0.8,
    verbose=False,
    n_ctx=2048,
    max_tokens=512,
)

Now that we have the model ready, let’s setup the prompt template like before, using the internal chat template.

from langchain_core.prompts import PromptTemplate

# Create a prompt template using OLMo's tokenizer chat template we saw in module 1.
prompt_template = PromptTemplate.from_template(
    template=olmo.client.metadata["tokenizer.chat_template"],
    template_format="jinja2",
    partial_variables={"add_generation_prompt": True, "eos_token": "<|endoftext|>"},
)

We again use the partial variables here to fill out the add_generation_prompt and eos_token fields. So that we’re left with just the messages input variables.

prompt_template.input_variables

['messages']

We have the prompt template ready, let’s move on to the next step, and create a prompt for the model. For simplicity of this tutorial, we’ll only use one message, user input to the model. This means we’ll only ask the model a single question at a time, rather than a series of questions that can feed of each other.

import textwrap  # a module to wrap text to make it more readable

Just like before, we’ll start by checking out what our full prompt text is going to look like. In this example, we’ve also used a handy built-in python module called textwrap to wrap the text to a certain width. We are using this to dedent the extra spaces to make it look cleaner.

# Test the prompt you want to send to OLMo.
question = "What is dark matter?"
input_content = textwrap.dedent(
    f"""\
    You are an astrophysics expert. Please answer the following question on astrophysics.
    Question: {question}
"""
)
input_messages = [
    {
        "role": "user",
        "content": input_content,
    }
]

full_prompt_text = prompt_template.format(messages=input_messages)

print(full_prompt_text)

<|endoftext|>

<|user|>
You are an astrophysics expert. Please answer the following question on astrophysics.
Question: What is dark matter?



<|assistant|>

Our prompt looks good. Let’s now make a chain and invoke it.

# Chain the prompt template and olmo
llm_chain = prompt_template | olmo

# Invoke the chain with a question and other parameters.
captured_answer = llm_chain.invoke({"messages": input_messages})

Dark matter is a theoretical particle or collection of particles that has been observed through its effect on visible matter, but not directly seen due to the absence of electromagnetic radiation emitted by it. In this sense, it is "dark" or invisible. However, various methods indicate the presence of this substance in the universe, with roughly 85% of all mass-energy in the observable universe being comprised of dark matter.

It was first introduced by scientists to explain the motion and structure of galaxies without considering the influence of normal (baryonic) matter. Dark matter's properties are yet to be precisely defined due to its non-observation, but it is hypothesized to have a large mass per unit volume. This substance interacts with the visible universe through gravitational forces only.

To date, scientists still seek evidence or direct detection of dark matter particles; however, different theoretical approaches and indirect methods provide compelling evidence for the existence of this fundamental component of the cosmos. Dark matter is essential to understand various cosmic phenomena like galaxy rotation curves, the structure formation in the universe, and the distribution of matter within our own Milky Way galaxy.

Great! At this point we have reviewed essentially module 1. But to ask different questions, we’ll need a way to pass in different questions to the chain. We know that we can just create new values for question, input_content, and input_messages variables, but that’s a lot of work and formatting to do every time we want to ask a new question. So what can we do?

Partial prompt#

We will now introduce a new concept called partial formatting. By using this feature, we can expand the input variables to be ones that we can easily change and pass in new values to. Essentially, we are creating a new prompt template from the underlying model template.

We’ve seen this feature in module 1 and above with the use of partial_variables in the model setup. This time, since we know that we’re only using one message, we can simplify the prompt template to take variables question and instruction.

First, let’s create a simple prompt template string that takes in the variables we want to pass in.

input_prompt_template = textwrap.dedent(
    """\
{instruction}

Question: {question}
"""
)

Notice that the above prompt template string is NOT an f-string, but rather the simple string, like the ones you’ve created in module 1.

Now that we have the prompt template string ready, let’s create a partial formatting from it. Remember that prompt_template is a String PromptTemplate object that contains the original jinja-2 template string with the variables add_generation_prompt and eos_token filled in. The only variable left is messages, which we will create a partial formatting with.

prompt_template

PromptTemplate(input_variables=['messages'], partial_variables={'add_generation_prompt': True, 'eos_token': '<|endoftext|>'}, template="{{ eos_token }}{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '<|user|>\n' + message['content'] }}\n{% elif message['role'] == 'assistant' %}\n{{ '<|assistant|>\n'  + message['content'] + eos_token }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '<|assistant|>' }}\n{% endif %}\n{% endfor %}", template_format='jinja2')

partial_prompt_template = prompt_template.partial(
    messages=[
        {
            "role": "user",
            "content": input_prompt_template,
        }
    ]
)

partial_prompt_template

PromptTemplate(input_variables=[], partial_variables={'add_generation_prompt': True, 'eos_token': '<|endoftext|>', 'messages': [{'role': 'user', 'content': '{instruction}\n\nQuestion: {question}\n'}]}, template="{{ eos_token }}{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '<|user|>\n' + message['content'] }}\n{% elif message['role'] == 'assistant' %}\n{{ '<|assistant|>\n'  + message['content'] + eos_token }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '<|assistant|>' }}\n{% endif %}\n{% endfor %}", template_format='jinja2')

As you can see above, the partial formatting is simply filling in the variables messages and now we’re left with no input_variables. So at this point, how can we create a new prompt template from this?

The answer is pretty straightforward. Let’s just call the .format and get the “final” prompt template string.

new_prompt_string = partial_prompt_template.format()

new_prompt_string

'<|endoftext|>\n\n<|user|>\n{instruction}\n\nQuestion: {question}\n\n\n\n<|assistant|>\n\n'

Now we have a simple prompt string that we can create a String PromptTemplate from.

new_prompt_template = PromptTemplate.from_template(new_prompt_string)

new_prompt_template

PromptTemplate(input_variables=['instruction', 'question'], template='<|endoftext|>\n\n<|user|>\n{instruction}\n\nQuestion: {question}\n\n\n\n<|assistant|>\n\n')

You can see now that the new prompt template takes in instruction and question. Let’s create a new chain and invoke it with this new prompt template.

Q&A Session with OLMo#

We’ll first create a single domain instruction, since we know that we’re asking questions about astrophysics.

domain_instruction = "You are an astrophysics expert. Please answer the following question on astrophysics."

question = "How many dimensions are there in the elemental abundances of stars?"

llm_chain = new_prompt_template.partial(instruction=domain_instruction) | olmo

llm_chain.invoke({"question": question})

The question seems to be asking about the number of dimensions involved in calculating the elemental abundances of stars. In astronomy, "elemental abundances" refer to the relative amounts of different elements present in a celestial body like a star, planet or asteroid. These abundance values are obtained by analyzing the composition of a sample and comparing it with known elemental ratios for each element.

The number of dimensions involved in this process is three: (1) The sample under study can have any shape (sphere, ellipsoid, etc.). In practice, it's typically a solid object or a piece(s) cut out from a bigger body like a star. (2) A sample has three spatial dimensions, namely x, y, and z, as we're generally talking about the composition of stars in three-dimensional space.
 (3) The elemental abundances refer to the number of atoms per unit volume of an element at each location within the sample. This means that the abundance of each element is a function of position in the sample (x, y, and z).

In summary, there are three dimensions: spatial (x, y, z), which have four components each; one time dimension (t) typically for measurements from various times during the star's life or from different locations on its surface; finally, the elemental abundances themselves, which are functions of these spatial and temporal coordinates.

It is important to mention that this example provides a simplified scenario for understanding elemental abundance values in stars. In reality, astronomers often encounter multi-phase systems like white dwarfs with distinct regions exhibiting distinct elemental compositions or other more complex cases. These additional complexities can further increase the number of dimensions involved in their analyses.

'The question seems to be asking about the number of dimensions involved in calculating the elemental abundances of stars. In astronomy, "elemental abundances" refer to the relative amounts of different elements present in a celestial body like a star, planet or asteroid. These abundance values are obtained by analyzing the composition of a sample and comparing it with known elemental ratios for each element.\n\nThe number of dimensions involved in this process is three: (1) The sample under study can have any shape (sphere, ellipsoid, etc.). In practice, it\'s typically a solid object or a piece(s) cut out from a bigger body like a star. (2) A sample has three spatial dimensions, namely x, y, and z, as we\'re generally talking about the composition of stars in three-dimensional space.\n (3) The elemental abundances refer to the number of atoms per unit volume of an element at each location within the sample. This means that the abundance of each element is a function of position in the sample (x, y, and z).\n\nIn summary, there are three dimensions: spatial (x, y, z), which have four components each; one time dimension (t) typically for measurements from various times during the star\'s life or from different locations on its surface; finally, the elemental abundances themselves, which are functions of these spatial and temporal coordinates.\n\nIt is important to mention that this example provides a simplified scenario for understanding elemental abundance values in stars. In reality, astronomers often encounter multi-phase systems like white dwarfs with distinct regions exhibiting distinct elemental compositions or other more complex cases. These additional complexities can further increase the number of dimensions involved in their analyses.'

Your Turn 😎#

You have two options:

  1. Use the questions provided at the beginning of this notebook and reuse the llm chain to ask questions about astrophysics.

  2. With the new prompt template new_prompt_template and the olmo model. Create a new chain with a different domain instruction, and ask questions about that domain.

Feel free to ask any questions you like, and see how OLMo responds to them! If you’re open to sharing, we’d love to hear about the questions you asked and the responses you received in the etherpad at https://etherpad.wikimedia.org/p/ipvVZZVxeP2JhpPxE4j6.

# Write your code here
Contents