GLiREL: Generalist Model for Relation Extraction

To prepare the TACRED dataset for few-shot relation classification tasks, you must clone the transformation utility, download the raw TACRED JSON files, convert them into relation-specific lists, and finally run the data transformation script to generate few-shot samples.

Follow these steps in order:

1. Clone the transformation repository: Obtain the Few_Shot_transformation_and_sampling utility.
2. Download and unzip TACRED: Get the raw JSON data from Stanford NLP.
3. Convert to relation-based lists: Use convert_dataset_to_list_by_relation.py to group instances by their relation type.
4. Transform data: Use data_transformation.py to create the final few-shot dataset based on a category split configuration.

# 1. Clone the utility
git clone https://github.com/ofersabo/Few_Shot_transformation_and_sampling.git
cd Few_Shot_transformation_and_sampling

# 2. Download TACRED
wget https://nlp.stanford.edu/~manning/xyzzy/tacred-with-json.zip 
unzip tacred-with-json.zip 

# 3. Convert dataset to lists by relation
mkdir -p tacred/data/instances_per_relation
python convert_dataset_to_list_by_relation.py --dataset tacred/data/json/train.json --output_file tacred/data/instances_per_relation/TACRED_train.json
python convert_dataset_to_list_by_relation.py --dataset tacred/data/json/dev.json --output_file tacred/data/instances_per_relation/TACRED_dev.json
python convert_dataset_to_list_by_relation.py --dataset tacred/data/json/test.json --output_file tacred/data/instances_per_relation/TACRED_test.json

# 4. Run data transformation
python data_transformation.py --train_data tacred/data/instances_per_relation/TACRED_train.json --dev_data tacred/data/instances_per_relation/TACRED_dev.json --test_data tacred/data/instances_per_relation/TACRED_test.json --fixed_categories_split categories_split.json --test_size 10 --output_dir ./data_few_shot

You can perform zero-shot relation extraction by using the GLiREL class. Load a pretrained model with GLiREL.from_pretrained and use predict_relations to find relations between entities.

Key Parameters for predict_relations:

• tokens: A list of tokens from the text.
• labels: A list of relation labels to search for.
• ner: A list of entity annotations in the format [start_index, end_index, entity_type, entity_text].
• threshold: Minimum confidence score (default 0.0).
• top_k: Number of top relations to return.

Note: Entity indices in the ner list are inclusive (e.g., [7, 7] refers to a single token at index 7).

from glirel import GLiREL
import spacy

model = GLiREL.from_pretrained("jackboyla/glirel-large-v0")

nlp = spacy.load('en_core_web_sm')

text = 'Derren Nesbitt had a history of being cast in "Doctor Who", having played villainous warlord Tegana in the 1964 First Doctor serial "Marco Polo".'
doc = nlp(text)
tokens = [token.text for token in doc]

labels = ['country of origin', 'licensed to broadcast to', 'father', 'followed by', 'characters']

# ner format: [start_index, end_index, type, text]
ner = [[26, 27, 'PERSON', 'Marco Polo'], [22, 23, 'Q2989412', 'First Doctor']]

relations = model.predict_relations(tokens, labels, threshold=0.0, ner=ner, top_k=1)

for item in relations:
    print(f"{item['head_text']} --> {item['label']} --> {item['tail_text']} | score: {item['score']}")

GLiREL can be added as a component to an existing spaCy pipeline. After adding the glirel pipe, you can pass tuples of (text, labels) to nlp.pipe to perform inference.

Steps:

1. Load a spaCy model.
2. Add the component using nlp.add_pipe("glirel", after="ner").
3. Pass data as tuples containing the text and the constrained labels dictionary.
4. Access results via the ._.relations attribute on the returned Doc objects.

import spacy
import glirel

nlp = spacy.load('en_core_web_sm')
nlp.add_pipe("glirel", after="ner")

text = "Apple Inc. was founded by Steve Jobs, Steve Wozniak, and Ronald Wayne in April 1976."
labels = {"glirel_labels": {
    'co-founder': {"allowed_head": ["PERSON"], "allowed_tail": ["ORG"]},
    'no relation': {},
}}

# Use nlp.pipe with as_tuples=True to pass labels
docs = list( nlp.pipe([(text, labels)], as_tuples=True) )
relations = docs[0][0]._.relations

for rel in relations:
    print(rel)

You can reduce noise by constraining which entity types are allowed to participate in specific relations. Pass a dictionary to the inference logic where keys are relation names and values define allowed_head and allowed_tail lists of entity labels.

Use constrain_relations_by_entity_type from glirel.modules.utils to apply these constraints to the model's output.

Constraint Schema:

labels_dict = {
    'relation_name': {
        'allowed_head': ['LABEL1', 'LABEL2'],
        'allowed_tail': ['LABEL3']
    }
}

from glirel.modules.utils import constrain_relations_by_entity_type

# Define constraints
labels_and_constraints = {
    'co-founder': {"allowed_head": ["PERSON"], "allowed_tail": ["ORG"]},
    'parent': {"allowed_head": ["PERSON"], "allowed_tail": ["PERSON"]}
}

# Apply to relations
relations = model.predict_relations(tokens, labels, threshold=0.0, ner=ner, top_k=1)
if isinstance(labels, dict): # If using the dict format
    relations = constrain_relations_by_entity_type(doc.ents, labels_and_constraints, relations)

To prepare a dataset for relation extraction, you can use spaCy to process text and extract entity pairs from sentences.

The create_data_items function iterates through a dataset, performs Named Entity Recognition (NER) using a spaCy model, and generates all possible non-identical entity pairs (e1 != e2) within sampled sentences. Each item includes the text, the head/tail pair, and a sentence_id to allow grouping related pairs later.

import spacy
import random
import tqdm

def create_data_items(nlp, dataset, n_sents=1):
    pair_items = []
    sent_id = 0
    for x in tqdm.tqdm(dataset):
        doc = nlp(x['text'])
        sents = list(doc.sents)
        sampled_sents = random.sample(sents, min(n_sents, len(sents)))
        for sent in sampled_sents:
            entities = [ent.text for ent in sent.ents]
            for e1 in entities:
                for e2 in entities:
                    if e1 != e2:
                        pair_items.append({
                            "relation": {"head": e1, "tail": e2}, 
                            "text": sent.text, 
                            "sentence_id": sent_id
                        })
            sent_id += 1
    return pair_items

You can use OpenAI's chat completion API to identify relationships between entity pairs in a text. By using Jinja2 templates, you can inject specific entity pairs and verbosity examples into prompts to guide the LLM.

To ensure reliable parsing, use response_format={ "type": "json_object" } in the OpenAI client call. This constrains the model to return a valid JSON object, which can then be parsed using json.loads().

Note that the relationship wording may change depending on which entity is the head and which is the tail.

import json
from openai import OpenAI
from jinja2 import Template

client = OpenAI()

def run_openai(rendered_prompt, system_prompt, model="gpt-3.5-turbo", max_tokens=512, temperature=0.5):
    response = client.chat.completions.create(
        model=model,
        response_format={ "type": "json_object"},
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": rendered_prompt}
        ],
        max_tokens=max_tokens,
        temperature=temperature
    )
    content = response.choices[0].message.content
    return json.loads(content)

You can use distilabel to build pipelines for generating synthetic datasets using LLMs. The workflow typically involves loading data (either from a dictionary or a Hugging Face Hub dataset), formatting it into prompts, and passing it through a TextGeneration step powered by an LLM like OpenAILLM.

from distilabel.llms import OpenAILLM
from distilabel.pipeline import Pipeline
from distilabel.steps import LoadDataFromDicts
from distilabel.steps.tasks import TextGeneration

# 1. Prepare your data
DATA = ["Your prompt template here".format(**item) for item in data]

# 2. Define the pipeline
with Pipeline(
    name="simple-text-generation-pipeline",
    description="A simple text generation pipeline",
) as pipeline:
    load_dataset = LoadDataFromDicts(
        name="load_data",
        data=[{"instruction": sample} for sample in DATA],
        batch_size=1,
    )
    generate_with_openai = TextGeneration(
        name="generate_with_gpt35", 
        llm=OpenAILLM(model="gpt-3.5-turbo")
    )
    load_dataset.connect(generate_with_openai)

# 3. Run the pipeline
distiset = pipeline.run(
    parameters={
        "generate_with_gpt35": {
            "llm": {
                "generation_kwargs": {
                    "temperature": 0.7,
                    "max_new_tokens": 512,
                }
            }
        },
    },
)

To prevent data leakage or mingling between zero-shot relation datasets and existing benchmarks (like WikiZSL, FewRel, or Re-DocRED), you can filter out relations that exist in both.

This process involves:

1. Identifying the intersection of relation labels between your zero-shot dataset and the benchmark datasets.
2. Using a filtering function to remove any relation in your dataset that matches an intersection label.
3. Removing entire items from the dataset if they no longer contain any valid relations after filtering.

Note: The filtering logic explicitly includes no_relation and no relation in the intersection labels to ensure they are treated as overlapping.

from tqdm import tqdm
import json

def remove_overlapping_relations(original_file_path, final_file_path, intersection_labels):
    # Ensure no-relation variants are treated as overlapping
    intersection_labels.add('no_relation')
    intersection_labels.add('no relation')

    skipped_items = 0
    skipped_relations = 0
    
    with open(original_file_path, 'r') as fr, open(final_file_path, 'w') as fw:
        for line in tqdm(fr):
            item = json.loads(line)
            relations = item['relations']
            new_relations = []
            
            for relation in relations:
                rel_text = relation['relation_text']
                if rel_text not in intersection_labels:
                    new_relations.append(relation)
                else:
                    skipped_relations += 1

            item['relations'] = new_relations

            # Only write the item if it still contains relations
            if len(new_relations) > 0:
                fw.write(json.dumps(item) + '\n')
            else:
                skipped_items += 1

    print(f'To not mingle with benchmark datasets, we skipped {skipped_items} items and {skipped_relations} relations')

The SpacyGLiRELWrapper class allows you to integrate GLiREL into a spaCy pipeline as a component. This enables zero-shot relation extraction on Doc objects that already contain named entities (doc.ents).

To use it, add the component to your pipeline using nlp.add_pipe("glirel", ...) and pass relation labels and constraints via the glirel_labels key in the doc._context during inference.

Workflow

1. Initialize: Create the wrapper with a pretrained GLiREL model.
2. Add to Pipeline: Use nlp.add_pipe("glirel", after="ner") (it is recommended to place it after an NER component).
3. Inference: Pass a tuple of (text, labels_dict) to nlp.pipe using as_tuples=True. The labels_dict must contain the key glirel_labels.
4. Access Results: Extracted relations are stored in the custom spaCy extension doc._.relations.

Relation Constraints

You can provide a dictionary for glirel_labels to constrain relations based on entity types using allowed_head and allowed_tail keys. This ensures the model only predicts relations that make sense for the specific entity types found.

# 1. Add the GLiREL component to the pipeline
nlp.add_pipe("glirel", after="ner")

# 2. Define labels and constraints
labels = {"glirel_labels": {
    'co-founder': {"allowed_head": ["PERSON"], "allowed_tail": ["ORG"]},
    'headquartered in': {"allowed_head": ["ORG"], "allowed_tail": ["LOC", "GPE", "FAC"]},
    'no relation': {},
}}

# 3. Run inference using nlp.pipe with as_tuples=True
text = "Apple Inc. was founded by Steve Jobs in Cupertino."
docs = list(nlp.pipe([(text, labels)], as_tuples=True))

# 4. Access the relations
relations = docs[0][0]._.relations
for rel in relations:
    print(rel)