Named Entity Recognition for Reddit Text Analysis

import spacy
from typing import List, Dict, Tuple
from collections import defaultdict

class RedditNERExtractor:
    """
    spaCy-based NER for Reddit text.

    Includes Reddit-specific preprocessing and
    custom entity normalization.
    """

    def __init__(
        self,
        model: str = "en_core_web_lg",
        entity_types: List[str] = None
    ):
        self.nlp = spacy.load(model)
        self.entity_types = entity_types or [
            'ORG', 'PRODUCT', 'PERSON', 'GPE', 'MONEY'
        ]

        # Entity normalization mappings
        self.normalizations = {
            # Stock tickers
            'aapl': 'Apple',
            'msft': 'Microsoft',
            'googl': 'Google',
            'amzn': 'Amazon',
            'nvda': 'NVIDIA',
            'tsla': 'Tesla',
            'meta': 'Meta',
            'nflx': 'Netflix',

            # Common variations
            'google': 'Google',
            'alphabet': 'Google',
            'iphone': 'iPhone',
            'macbook': 'MacBook',
            'openai': 'OpenAI',
            'chatgpt': 'ChatGPT',
            'gpt-4': 'GPT-4',
            'gpt4': 'GPT-4',
        }

    def normalize_entity(self, text: str) -> str:
        """Normalize entity text to canonical form."""
        normalized = self.normalizations.get(text.lower())
        if normalized:
            return normalized
        return text

    def extract(self, text: str) -> Dict[str, List[Dict]]:
        """
        Extract entities from text.

        Returns dict mapping entity type to list of entities.
        """
        doc = self.nlp(text)

        entities = defaultdict(list)

        for ent in doc.ents:
            if ent.label_ in self.entity_types:
                normalized = self.normalize_entity(ent.text)

                entities[ent.label_].append({
                    'text': ent.text,
                    'normalized': normalized,
                    'start': ent.start_char,
                    'end': ent.end_char,
                    'label': ent.label_
                })

        return dict(entities)

    def extract_batch(
        self,
        texts: List[str],
        batch_size: int = 100
    ) -> List[Dict[str, List[Dict]]]:
        """Extract entities from multiple texts efficiently."""
        results = []

        # Use spaCy's pipe for batch processing
        for doc in self.nlp.pipe(texts, batch_size=batch_size):
            entities = defaultdict(list)

            for ent in doc.ents:
                if ent.label_ in self.entity_types:
                    normalized = self.normalize_entity(ent.text)
                    entities[ent.label_].append({
                        'text': ent.text,
                        'normalized': normalized,
                        'label': ent.label_
                    })

            results.append(dict(entities))

        return results

    def get_entity_counts(
        self,
        texts: List[str]
    ) -> Dict[str, Dict[str, int]]:
        """Count entity occurrences across multiple texts."""
        counts = defaultdict(lambda: defaultdict(int))

        for result in self.extract_batch(texts):
            for entity_type, entities in result.items():
                for entity in entities:
                    counts[entity_type][entity['normalized']] += 1

        return {k: dict(v) for k, v in counts.items()}


# Usage
extractor = RedditNERExtractor()

text = """
    Just sold my TSLA shares and bought AAPL instead.
    Tim Cook's strategy with the iPhone 15 Pro is brilliant.
    Microsoft's Copilot is giving OpenAI a run for their money.
"""

entities = extractor.extract(text)
for entity_type, ents in entities.items():
    print(f"\n{entity_type}:")
    for e in ents:
        print(f"  {e['text']} -> {e['normalized']}")

from transformers import AutoTokenizer, AutoModelForTokenClassification
from transformers import pipeline
from typing import List, Dict
import torch

class TransformerNER:
    """
    Transformer-based NER for high-accuracy extraction.

    Uses fine-tuned BERT/RoBERTa models for better
    handling of informal text.
    """

    def __init__(
        self,
        model_name: str = "dslim/bert-base-NER",
        device: int = -1  # -1 for CPU, 0+ for GPU
    ):
        self.ner_pipeline = pipeline(
            "ner",
            model=model_name,
            tokenizer=model_name,
            aggregation_strategy="simple",
            device=device
        )

        # Map model labels to standard types
        self.label_map = {
            'B-PER': 'PERSON', 'I-PER': 'PERSON',
            'B-ORG': 'ORG', 'I-ORG': 'ORG',
            'B-LOC': 'GPE', 'I-LOC': 'GPE',
            'B-MISC': 'MISC', 'I-MISC': 'MISC',
            'PER': 'PERSON',
            'ORG': 'ORG',
            'LOC': 'GPE',
        }

    def extract(
        self,
        text: str,
        min_score: float = 0.7
    ) -> List[Dict]:
        """
        Extract entities with confidence filtering.

        Args:
            text: Input text
            min_score: Minimum confidence threshold
        """
        # Truncate very long texts
        if len(text) > 5000:
            text = text[:5000]

        results = self.ner_pipeline(text)

        entities = []
        for r in results:
            if r['score'] >= min_score:
                label = self.label_map.get(r['entity_group'], r['entity_group'])
                entities.append({
                    'text': r['word'],
                    'label': label,
                    'score': r['score'],
                    'start': r['start'],
                    'end': r['end']
                })

        return entities

    def extract_batch(
        self,
        texts: List[str],
        batch_size: int = 16
    ) -> List[List[Dict]]:
        """Batch extraction for efficiency."""
        all_results = []

        for i in range(0, len(texts), batch_size):
            batch = texts[i:i + batch_size]
            batch_results = self.ner_pipeline(batch)

            for results in batch_results:
                entities = []
                for r in results:
                    if r['score'] >= 0.7:
                        label = self.label_map.get(
                            r['entity_group'],
                            r['entity_group']
                        )
                        entities.append({
                            'text': r['word'],
                            'label': label,
                            'score': r['score']
                        })
                all_results.append(entities)

        return all_results


# Usage
ner = TransformerNER(device=0)  # GPU

text = "Elon Musk announced that Tesla will partner with Microsoft on AI."
entities = ner.extract(text)

for e in entities:
    print(f"{e['text']} ({e['label']}): {e['score']:.2f}")

from transformers import (
    AutoModelForTokenClassification,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForTokenClassification
)
from datasets import Dataset
import numpy as np
from seqeval.metrics import f1_score, precision_score, recall_score

class RedditNERTrainer:
    """Fine-tune NER model on Reddit data."""

    def __init__(
        self,
        base_model: str = "bert-base-uncased",
        labels: List[str] = None
    ):
        self.labels = labels or [
            'O',
            'B-ORG', 'I-ORG',
            'B-PRODUCT', 'I-PRODUCT',
            'B-PERSON', 'I-PERSON',
            'B-MONEY', 'I-MONEY',
        ]

        self.label2id = {l: i for i, l in enumerate(self.labels)}
        self.id2label = {i: l for l, i in self.label2id.items()}

        self.tokenizer = AutoTokenizer.from_pretrained(base_model)
        self.model = AutoModelForTokenClassification.from_pretrained(
            base_model,
            num_labels=len(self.labels),
            id2label=self.id2label,
            label2id=self.label2id
        )

    def tokenize_and_align(self, examples):
        """Tokenize and align labels with subwords."""
        tokenized = self.tokenizer(
            examples['tokens'],
            truncation=True,
            is_split_into_words=True,
            padding=True
        )

        labels = []
        for i, label in enumerate(examples['ner_tags']):
            word_ids = tokenized.word_ids(batch_index=i)
            label_ids = []
            previous_word_idx = None

            for word_idx in word_ids:
                if word_idx is None:
                    label_ids.append(-100)
                elif word_idx != previous_word_idx:
                    label_ids.append(self.label2id[label[word_idx]])
                else:
                    # For subwords, use I- prefix if B-
                    curr_label = label[word_idx]
                    if curr_label.startswith('B-'):
                        label_ids.append(self.label2id['I-' + curr_label[2:]])
                    else:
                        label_ids.append(self.label2id[curr_label])

                previous_word_idx = word_idx

            labels.append(label_ids)

        tokenized['labels'] = labels
        return tokenized

    def compute_metrics(self, eval_preds):
        """Compute NER metrics."""
        predictions, labels = eval_preds
        predictions = np.argmax(predictions, axis=2)

        true_labels = []
        true_predictions = []

        for prediction, label in zip(predictions, labels):
            true_l = []
            true_p = []
            for p, l in zip(prediction, label):
                if l != -100:
                    true_l.append(self.id2label[l])
                    true_p.append(self.id2label[p])
            true_labels.append(true_l)
            true_predictions.append(true_p)

        return {
            'precision': precision_score(true_labels, true_predictions),
            'recall': recall_score(true_labels, true_predictions),
            'f1': f1_score(true_labels, true_predictions)
        }

    def train(
        self,
        train_data: List[Dict],
        val_data: List[Dict],
        output_dir: str = "./reddit-ner-model",
        epochs: int = 3
    ):
        """Train the NER model."""
        # Prepare datasets
        train_dataset = Dataset.from_list(train_data)
        val_dataset = Dataset.from_list(val_data)

        train_tokenized = train_dataset.map(
            self.tokenize_and_align,
            batched=True
        )
        val_tokenized = val_dataset.map(
            self.tokenize_and_align,
            batched=True
        )

        training_args = TrainingArguments(
            output_dir=output_dir,
            num_train_epochs=epochs,
            per_device_train_batch_size=16,
            per_device_eval_batch_size=16,
            warmup_ratio=0.1,
            weight_decay=0.01,
            logging_steps=100,
            eval_strategy="epoch",
            save_strategy="epoch",
            load_best_model_at_end=True,
            metric_for_best_model="f1",
        )

        data_collator = DataCollatorForTokenClassification(self.tokenizer)

        trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=train_tokenized,
            eval_dataset=val_tokenized,
            tokenizer=self.tokenizer,
            data_collator=data_collator,
            compute_metrics=self.compute_metrics
        )

        trainer.train()
        trainer.save_model(output_dir)

        return trainer