# Inertial Confinement Fusion Forecasting via Large Language Models

This repository is the official implementation of Inertial Confinement Fusion Forecasting via Large Language Models. 

![Main Figure: LPI-LLM framework.](images/main-fig_v2.png)
<p align="center">
 Main Figure: LPI-LLM framework. (a) Fusion-specific prompts structure domain textual prompts with context, task and input descriptions (see §2.2.1 for details). (b) Signal-digesting channels comprise a pre-trained temporal encoder to extract time-series features of laser signals, and a spatial encoder to encode critical landscapes of the inputs (see §2.2.2 for details). For simplicity, we skip some architectural modules. We provide implementation details in appendix (see §S1).
</p>

## Requirements

To install requirements:

```setup
conda env create -f environment.yml
```

Note that depending on your network environment, this command may take quite a while to execute.

Whether running training or evaluation code, please ensure that internet access - especially to huggingface - is smooth. Our code relies on weights hosted on huggingface. Storage that can at least hold the weights of LLM is also required.

## Dataset and Weights

One of the key contributions of this study, the LPI4AI dataset, will be made publicly available upon the acceptance of this paper. In the interim, we have included the test split of the LPI4AI dataset (located at dataset/testset_cleaned_sampled.pickle) in this repository, along with the fine-tuned weights of our model necessary to replicate our experimental findings. 

At this moment we are temporarily hosting our model weights on Dropbox, please use [this link](https://www.dropbox.com/scl/fi/xr2454n6okhkggn71j2ok/checkpoint?rlkey=u9ln7u5iyfutkjx77pf97gfvq&e=1&st=1dc26z8e&dl=0) to download the weights and place it in the ./weights folder in order to reproduce the results of the experiment.

## Training

To train the model(s) in the paper, run this command:

```train
bash scripts/train.sh
```

Please adjust hyper-parameters such as the number of GPUs in this script according to your actual setup.

## Evaluation

To evaluate my model on LPI4AI dataset, run:

```eval
bash scripts/eval.sh
```

## Results

LPI-LLM achieves the following performance on test split of LPI4AI:

| Model name         | CAE             | Top 1 MAE      | Top 5 MAE      |
| ------------------ |---------------- | -------------- | -------------- |
| LPI-LLM         |     1.90±0.33    |      0.14±0.01      |      0.11±0.01      |