README.md

---
license: apache-2.0
base_model: Qwen/Qwen3-8B
tags:
- coreference-resolution
- ner
- entity-extraction
- qwen3
- lora
- peft
datasets:
- wjbmattingly/synthetic-coref
language:
- en
pipeline_tag: text-generation
---

# Qwen3-8B-Coref-NER

A fine-tuned [Qwen3-8B](Qwen/Qwen3-8B) model for **coreference resolution** and **named entity recognition**. 

This model resolves pronouns and other referring expressions by replacing them with the full entity names, while also tracking entity mentions and their variants.

## Model Description

- **Base Model:** [Qwen/Qwen3-8B](https://huggingface.co/Qwen/Qwen3-8B)
- **Training Dataset:** [wjbmattingly/synthetic-coref](https://huggingface.co/datasets/wjbmattingly/synthetic-coref)
- **Task:** Coreference Resolution + Entity Tracking
- **Method:** LoRA (Low-Rank Adaptation)

## Usage

### Installation

```bash
pip install transformers peft torch
```

### Quick Start

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# Load base model and tokenizer
base_model = "Qwen/Qwen3-8B"
tokenizer = AutoTokenizer.from_pretrained(base_model)
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

# Load LoRA adapter
model = PeftModel.from_pretrained(model, "wjbmattingly/Qwen3-8B-Coref-NER")

# Sample text
text = """Alcuin of York was an Anglo-Latin scholar and teacher. He was born around 735 and became the student of Archbishop Ecgbert at York. At the invitation of Charlemagne, he became a leading scholar at the Carolingian court.

In this role as adviser, he took issue with the emperor's policy of forcing pagans to be baptised on pain of death. His arguments seem to have prevailed – Charlemagne abolished the death penalty for paganism in 797."""

# Create prompt
prompt = "Resolve all pronouns in this text, replacing them with the full entity names. Also identify any entity references you find.\n\n" + text

messages = [{"role": "user", "content": prompt}]
input_text = tokenizer.apply_chat_template(
    messages, 
    tokenize=False, 
    add_generation_prompt=True,
    enable_thinking=False  # Disable thinking mode
)

# Generate
inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=2048,
        do_sample=False,
        pad_token_id=tokenizer.pad_token_id,
    )

response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
print(response)
```

### Paragraph-by-Paragraph Processing with Entity Tracking

For longer documents, process paragraph by paragraph while tracking entities:

```python
import torch
import re
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

def parse_entity_mappings(response):
    """Parse model response to extract resolved text and entity mappings."""
    if "NEW ENTITY MAPPINGS:" in response:
        parts = response.split("NEW ENTITY MAPPINGS:")
        resolved_text = parts[0].strip()
        mappings_text = parts[1].strip() if len(parts) > 1 else ""
        
        entities = {}
        for line in mappings_text.split("\n"):
            line = line.strip()
            if line.startswith("-"):
                match = re.match(r'-\s*([^:]+):\s*\[([^\]]*)\]', line)
                if match:
                    entity_name = match.group(1).strip()
                    variants = re.findall(r'"([^"]*)"', match.group(2))
                    if variants:
                        entities[entity_name] = variants
        return resolved_text, entities
    return response.strip(), {}

def format_entities_for_prompt(entities):
    """Format known entities for the prompt."""
    lines = ["Entities and their possible references:"]
    for entity_name, variants in entities.items():
        variants_str = ", ".join(f'"{v}"' for v in variants)
        lines.append(f"- {entity_name}: [{variants_str}]")
    return "\n".join(lines)

# Load model
base_model = "Qwen/Qwen3-8B"
tokenizer = AutoTokenizer.from_pretrained(base_model)
model = AutoModelForCausalLM.from_pretrained(base_model, torch_dtype=torch.bfloat16, device_map="auto")
model = PeftModel.from_pretrained(model, "wjbmattingly/Qwen3-8B-Coref-NER")

# Your document
text = """Alcuin of York was an Anglo-Latin scholar and teacher. He was born around 735 and became the student of Archbishop Ecgbert at York. At the invitation of Charlemagne, he became a leading scholar at the Carolingian court.

In this role as adviser, he took issue with the emperor's policy of forcing pagans to be baptised on pain of death. His arguments seem to have prevailed – Charlemagne abolished the death penalty for paganism in 797."""

# Split into paragraphs
paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
resolved_paragraphs = []
cumulative_entities = {}

for i, paragraph in enumerate(paragraphs):
    print(f"Processing paragraph {i+1}/{len(paragraphs)}...")
    
    # Build prompt
    if i == 0:
        prompt = f"Resolve all pronouns in this text, replacing them with the full entity names. Also identify any entity references you find.\n\n{paragraph}"
    else:
        context = "\n\n".join(resolved_paragraphs[max(0, i-2):i])
        if cumulative_entities:
            known_str = format_entities_for_prompt(cumulative_entities)
            prompt = f"Known {known_str}\n\nGiven this context of preceding text (already resolved):\n\n{context}\n\nResolve all pronouns in this paragraph using the known entities. Also identify any NEW entity references:\n\n{paragraph}"
        else:
            prompt = f"Given this context of preceding text (already resolved):\n\n{context}\n\nResolve all pronouns in this paragraph. Also identify any NEW entity references:\n\n{paragraph}"
    
    messages = [{"role": "user", "content": prompt}]
    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=2048, do_sample=False, pad_token_id=tokenizer.pad_token_id)
    
    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()
    
    # Parse response
    resolved, new_entities = parse_entity_mappings(response)
    resolved_paragraphs.append(resolved)
    
    # Update cumulative entities
    for entity_name, variants in new_entities.items():
        if entity_name not in cumulative_entities:
            cumulative_entities[entity_name] = set()
        cumulative_entities[entity_name].update(variants)
    
    if new_entities:
        print(f"  New entities found: {new_entities}")

# Final output
print("\n" + "="*50)
print("RESOLVED TEXT:")
print("="*50)
print("\n\n".join(resolved_paragraphs))

print("\n" + "="*50)
print("ALL ENTITIES:")
print("="*50)
for entity, variants in cumulative_entities.items():
    print(f"  {entity}: {list(variants)}")
```

## Sample Output

**Input:**
```
Alcuin of York was an Anglo-Latin scholar and teacher. He was born around 735 and became the student of Archbishop Ecgbert at York. At the invitation of Charlemagne, he became a leading scholar at the Carolingian court.

In this role as adviser, he took issue with the emperor's policy of forcing pagans to be baptised on pain of death. His arguments seem to have prevailed – Charlemagne abolished the death penalty for paganism in 797.
```

**Output:**
```
Alcuin of York was an Anglo-Latin scholar and teacher. Alcuin of York was born around 735 and became the student of Archbishop Ecgbert at York. At the invitation of Charlemagne, Alcuin of York became a leading scholar at the Carolingian court.

In Alcuin of York's role as adviser, Alcuin of York took issue with Charlemagne's policy of forcing pagans to be baptised on pain of death. Alcuin of York's arguments seem to have prevailed – Charlemagne abolished the death penalty for paganism in 797.

NEW ENTITY MAPPINGS:
- Alcuin of York: ["He", "his", "he"]
- Charlemagne: ["the emperor"]
```

## Training Details

This model was trained using:
- **LoRA rank:** 16
- **LoRA alpha:** 32
- **Target modules:** q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj
- **Training mode:** Paragraph-by-paragraph with progressive entity tracking

## Citation

If you use this model, please cite the training dataset and base model.

## License

Apache 2.0