Table of Contents

• Introduction
• Demonstration
• Procedure of Data Construction
  • Mainframe-Training
  • Mainframe-Instruct
• Model Download
• Evaluation Results
• Usage
  • Fine-tune XMAiNframe
  • Inference
• License
• Acknowledgements
• Contact Us
• Citation Information

Introduction

We are introducing XMAiNframe, a state-of-the-art large language model (LLM) specifically designed with knowledge of mainframe legacy systems and COBOL codebases. XMAiNframe is built on top of DeepSeek-Coder 7B and is available with 7B and 10.5B parameters. Additionally, we present MainframeBench, a comprehensive benchmark for assessing mainframe knowledge, including multiple-choice questions, question answering, and COBOL code summarization. Our empirical evaluations demonstrate that XMAiNframe consistently outperforms existing state-of-the-art LLMs across these tasks. Specifically, XMAiNframe achieves 30% higher accuracy than DeepSeek-Coder on multiple-choice questions, doubles the BLEU score of Mixtral-Instruct 8x7B on question answering, and scores six times higher than GPT-3.5 on COBOL summarization. Our work highlights the potential of XMAiNframe to drive significant advancements in managing and modernizing legacy systems, thereby enhancing productivity and saving time for software developers.

Demonstration

In this section, we demonstrate the capabilities of XMAiNframe by comparing it with the leading language model, DeepSeek-Coder-7B. We evaluate the performance of each model by showcasing their responses to a series of realistic questions related to mainframe knowledge. The images below illustrate how each model handles identical prompts. As shown, the responses generated by XMAiNframe are not only accurate but also more detailed and comprehensive compared to those from the base model, DeepSeek-Coder-7B. This makes XMAiNframe particularly valuable for developers seeking a reliable and thorough AI assistant in the mainframe environment.

Procedure of Data Construction

Mainframe-Training

We utilized two different sources: using the GitHub API to collect COBOL projects hosted on GitHub and gathering online document data relevant to mainframes. In total, Mainframe-Training Dataset consists of 236 million tokens from documents about the mainframe technology and COBOL constructs. In the pre-training process, we combined our Mainframe-Training Dataset with SlimOrca-Dedup to enrich the model’s mainframe knowledge while retaining its general capabilities.

Mainframe-Instruct

Mainframe-Instruct is a high-quality synthetic dataset created through 5 steps:

• Step 1: 300 seed data instances about Mainframe and COBOL are gathered and annotated by our domain experts.

• Step 2: Using popular LLMs to enrich Mainframe-Instruct from seed data.

• Step 3: Utilizing GPT-4 as an evaluator to judge model responses, scoring the outputs and ranking responses in a pairwise manner.

• Step 4: Filtering and manually checking.

• Step 5: Dividing Mainframe-Instruct into three tasks: Multiple Choice Questions, Question Answering, and COBOL summarization.

Below are the statistics of Mainframe-Instruct Dataset:

MainframeBench, our benchmark for mainframe knowledge, is the testing set in Mainframe-Instruct Dataset. This benchmark is used to evaluate our LLMs with others which is now available at Huggingface datasets.

from datasets import load_dataset

# Load each sub-set in MainframeBench
QA_set = load_dataset("Fsoft-AIC/MainframeBench", 'question_answering')
MC_set = load_dataset("Fsoft-AIC/MainframeBench", 'multiple_choice_question')
Summarization_set = load_dataset("Fsoft-AIC/MainframeBench", 'COBOL_code_summarization')

Model Download

We release XMAiNframe with 7B and 10.5B parameters, including base and instruct models, to the public. XMAiNframe 10.5B is expanded from DeepSeek-Coder 7B by the depth up-scaling method without introducing additional modules or dynamic expert selection methods.

Evaluation Results

Multiple Choice Question Task

Question Answering Task

COBOL Code Summarization

For more evaluation details and settings, please check our paper.

Usage

Fine-tune XMAiNframe

To run the code in this project, first, create a Python virtual environment using e.g. Conda:

conda create -n xmainframe python=3.10 && conda activate xmainframe

You can then install the remaining package dependencies as follows:

git clone https://github.com/FSoft-AI4Code/XMainframe.git
cd XMainframe
pip install -r requirements.txt

You can now check out the scripts and recipes directories for instructions on how to fine-tune our model 🪁!

Inference

Here is a code snippet with apply_chat_template to show you how to load the tokenizer and model and how to generate content.

from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("Fsoft-AIC/XMAiNframe-instruct-7b")
model = AutoModelForCausalLM.from_pretrained("Fsoft-AIC/XMAiNframe-instruct-7b")
messages=[
    {'from':'system', 'value': "You are a helpful assistant"},
    {'from': 'human', 'value': 'What is the future of Mainframe?'}
]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
 
outputs = model.generate(inputs, max_new_tokens=512, do_sample=False, top_k=50, top_p=0.95, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)
print(tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True))

License

This code repository is licensed under the MIT License

Acknowledgements

This codebase is adapted from:

• alignment-handbook

Contact us

If you have any questions, comments or suggestions, please do not hesitate to contact us.

• Website: fpt-aicenter
• Email: [email protected]

Citation Information

More details can be found in our technical report.

If you're using XMAiNframe, please cite using this BibTeX:

@article{dau2024xmainframe,
  title={XMainframe: A Large Language Model for Mainframe Modernization},
  author={Dau, Anh TV and Dao, Hieu Trung and Nguyen, Anh Tuan and Tran, Hieu Trung and Nguyen, Phong X and Bui, Nghi DQ},
  journal={arXiv preprint arXiv:2408.04660},
  year={2024}
}