re-docred-mistral-7b-finetuned

docre-mistral-7b-v0.3

This model is fine-tuned on the RE-DocRED dataset for document-level relation extraction. It can extract relationships between entities in documents and is designed for building knowledge graphs from text.

Model Details

• Base Model: Mistral-7B-v0.3
• Task: Relation Extraction
• Dataset: RE-DocRED
• Number of Relation Types: 96

Performance

Supported Relations

The model can classify 96 relation types from the wikidata relationship set

Usage

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load model and tokenizer
model_name = "docre-mistral-7b-v0.3"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Example usage
text = "Your document text here."
head_entity = "Entity 1"
tail_entity = "Entity 2"

# Format input
input_text = f"{text} [SEP] {head_entity} [SEP] {tail_entity}"

# Tokenize and predict
inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
    outputs = model(**inputs)
    predicted_class = torch.argmax(outputs.logits, dim=-1)

# Get relation label
relation = model.config.id2label[predicted_class.item()]
print(f"Predicted relation: {relation}")

Training Data

This model was trained on the RE-DocRED dataset.

Citation

If you use this model, please cite:

@inproceedings{tan-etal-2022-revisiting,
    title = "Revisiting DocRED -- Addressing the False Negative Problem in Relation Extraction",
    author = "Tan, Qingyu and He, Ruidan and Bing, Lidong and Ng, Hwee Tou",
    booktitle = "Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing",
    year = "2022"
}

Disclaimer

This model is intended for research and educational purposes. Please ensure appropriate use and consider potential biases in the training data.