This repository is designed for editing input text using a Language Model. It allows users to apply various editing prompts and various models defined in configuration files to modify the input text.
Currently, the editing prompts are written to recognize and replace name entities such as names or locations from free text and replaces all occurrences with a placeholder defined in the prompt config file.
This project aims to stay model agnostic (i.e. it can be used with a model of the user's choice) and therefore avoid any vendor lock-in.
This software functions in a way like a smart editor. E.g. it can anonymize names in a text or exchange name entities for a batch of emails.
To use the AI-NER, follow these steps:
- Clone the repository:
git clone https://github.com/jWinman91/ai-extractor.git
cd ai-extractor- Install the required dependencies:
pip install -r requirements.txt- Download a model of your choice into
models. I recommend the following models from Hugging Face for German text:
Each model can be downloaded by using wget, e.g.:
wget https://huggingface.co/TheBloke/SauerkrautLM-7B-v1-mistral-GGUF/resolve/main/sauerkrautlm-7b-v1-mistral.Q4_0.gguf
In order to use this repository, several configuration need to be set for the model as well as the NER tasks to extract name entities.
These can be set in two types of configuration files, config_model and config_prompt.
config_modelsets all configurations necessary for the respective model.config_promptsets the configurations for the NER tasks (e.g. which model to choose and with what to replace the identified name entity).
TODO
After setting the configuration and downloading one (or more) of the models, you can simply use AI-NER by running:
python main.py $PATH_TO_INPUT $PATH_TO_OUTPUTAn example text file is added in data/input/email_example_de.txt, which is a self-written email in German.
There are also pre-defined config_model and config_prompt files.
By running AI-NER with the anonymize_example_email.yaml prompt configuration and german_mistral.yaml as well as the flair.yaml model configuration,
we can now anonymize certain entities in the example email.
Below are an image of before and after running python main.py on the email using the anonymize_emails-NER.yaml config file.
This project is licensed under the MIT License - see the LICENSE file for details.
- NLTK - Natural Language Toolkit used for sentence tokenization.
- Hugging Face - Framework for working with state-of-the-art natural language processing models.