Fine-Tuning Mistral 7B for Named Entity Recognition (NER)

Walid Amamou
UBIAI NLP
Published in
7 min readApr 18, 2024

--

In the field of natural language processing (NLP), Named Entity Recognition (NER) is recognized as a crucial task with wide-ranging applications, including information extraction and question answering systems.

The introduction of Mistral 7B, an innovative open-source large language model developed by Mistral AI, has brought about a significant transformation in the NLP landscape. This article aims to delve into the capabilities of Mistral 7B for NER tasks, focusing on the process of fine-tuning this advanced model to achieve superior performance in entity recognition.

Introduction to Named Entity Recognition

Named Entity Recognition (NER) is a fundamental task in NLP, focused on identifying and categorizing named entities within text, such as person names, locations, organizations, dates, and more. The precise identification of named entities is critical for numerous downstream applications, including information retrieval, sentiment analysis, and knowledge graph construction.

Traditionally, NER systems relied on manual rule creation and feature engineering techniques, which often lacked scalability and robustness across different domains.

However, the emergence of deep learning and pretrained language models like Mistral 7B has revolutionized the field. These models offer a data-driven approach that harnesses extensive text data, leading to improved performance and scalability in NER tasks.

Forging Ahead: Mistral 7B in Comparison to Other Open Source LLMs

Mistral 7B stands out as a leading contender, offering a unique blend of performance, efficiency, and accessibility. Distinguishing itself from its counterparts, Mistral 7B incorporates innovative architectural elements such as grouped-query attention and sliding window attention, elevating both inference speed and sequence handling capabilities.

These advancements not only differentiate Mistral 7B but also underscore its exceptional adaptability and resilience across a wide array of natural language processing (NLP) tasks. Moreover, Mistral 7B’s Apache 2.0 licensing and community-driven development approach cultivate an atmosphere of collaboration and innovation, cementing its status as a preferred choice among researchers and practitioners seeking state-of-the-art language understanding models.

Unveiling Mistral 7B: A Game-Changer in NLP

Mistral 7B emerges as an innovative powerhouse in the domain of language models, featuring an impressive 7 billion parameters meticulously crafted to address the intricacies of natural language.

Engineered with state-of-the-art functionalities like grouped-query attention (GQA) and sliding window attention (SWA), Mistral 7B demonstrates unmatched performance and efficiency, establishing itself as a formidable choice for NLP tasks, including Named Entity Recognition.

The incorporation of GQA allows Mistral 7B to accelerate inference, enabling real-time processing of textual data — an essential capability for NER tasks operating in dynamic environments. Additionally, SWA empowers Mistral 7B to efficiently handle sequences of varying lengths, ensuring robust performance across a diverse range of text inputs.

Fine-Tuning Mistral 7B for NER: Unlocking Its Full Potential

Although Mistral 7B exhibits impressive capabilities in its pre-trained state, its full potential is realized through the fine-tuning process. During fine-tuning, the model’s parameters are adjusted to suit the intricacies of particular tasks and datasets. Fine-tuning Mistral 7B for Named Entity Recognition (NER) tasks entails additional training on annotated data, enabling the model to grasp domain-specific patterns and enhance its accuracy in identifying entities.

Install the necessary Python packages and libraries required for fine-tuning Mistral

Data annotation and structure using UBIAI tools

In fine-tuning large language models like Mistral 7B, the quality and relevance of the data are critical factors that significantly influence the model’s performance and effectiveness. Utilizing high-quality, well-annotated data ensures that the model learns from accurate and representative examples, thereby improving its ability to generalize and generate reliable outputs.

For this tutorial, the data is sourced from UBIAI, a platform known for its robust annotation tools and comprehensive NLP capabilities. Leveraging UBIAI’s annotation capabilities ensures that the data used for fine-tuning is meticulously labeled and tailored to the specific task at hand, maximizing the effectiveness of the training process.

This code reads JSON data from a file containing invoice information. It organizes the data into a structured format, ensuring there are no duplicate annotations. The processed results are stored for further analysis and fine-tuning Mistral for Named Entity Recognition (NER) tasks.

1371

Distinct Labels: {‘PERSON’, ‘PRODUCT’, ‘FACILITY’, ‘MONEY’, ‘EVENT’, ‘DATE’, ‘ORGANIZATION’, ‘LOCATION’}

The dataset comprises documents containing names, texts, and annotations. These annotations encompass various entities like persons, organizations, locations, facilities, events, among others, offering structured data for Named Entity Recognition (NER) tasks.

The dataset is now augmented with a comprehensive prompt delineating instructions for entity extraction tasks. Each instruction delineates precise guidelines for extracting entities across various labels, including DATE, PERSON, ORGANIZATION, LOCATION, FACILITY, EVENT, MONEY, and PRODUCT. The prompt underscores the importance of accuracy and relevance in responses, offering guidelines to ensure consistency and comprehensiveness in the extracted information.

This structured prompt facilitates streamlined processing of the data for entity recognition tasks, aligning seamlessly with the objectives of the NER (Named Entity Recognition) process.

This process involves converting a subset of data into string representations, appending them to a list, and saving them to a CSV file. Subsequently, the CSV file is read into a Pandas DataFrame and transformed into a Hugging Face Datasets object for further processing.

the Mistral 7B base model is loaded using the BitsAndBytesConfig configuration, which enables quantization for efficient memory usage and faster inference. Additionally, the tokenizer associated with Mistral 7B is loaded, ensuring proper padding and end-of-sequence token settings for text processing tasks.

This section of the code sets up monitoring for the fine-tuning process of Mistral 7B using the Weights & Biases (W&B) platform. By authenticating with W&B using the provided API key, the training progress and performance metrics are continuously monitored in real-time. The `wandb.init()` function initiates a new run within the designated project (“Fine-tuning Mistral 7B UBIAI”), with the job type specified as “training”. Enabling anonymous access ensures the security of the training data while facilitating collaborative monitoring and analysis.

initializing hyperparameters for the training process, defining settings such as output directory, number of epochs, batch size, learning rate, and logging frequency. Additionally, it configures the SFTTrainer, specifying the model, dataset, tokenizer, and other relevant parameters essential for fine-tuning Mistral 7B for the given task.

Let’s Try Our Model: Testing Mistral 7B with User Prompts

The provided function stream takes a user prompt as input and generates text based on that prompt using the Mistral 7B model. It utilizes the model’s generate method to produce text, using the user’s prompt as context. The generated text is then returned as output.

Response:

Conclusion

Fine-tuning Mistral 7B for Named Entity Recognition represents a transformative approach to entity recognition tasks, offering enhanced performance, adaptability, and efficiency. As organizations continue to leverage NER for information extraction, knowledge discovery, and beyond, the capabilities unlocked by Mistral 7B pave the way for groundbreaking advancements in NLP. By harnessing the power of Mistral 7B and fine-tuning it for NER, practitioners can unlock new possibilities in understanding and processing textual data, driving innovation across diverse domains.

Credits of this article goes to: Ilyes Ben Khalifa

--

--

Walid Amamou
UBIAI NLP

Founder of UBIAI, annotation tool for NLP applications| PhD in Physics.