| import json |
| from pathlib import Path |
| from src.schemas.labels import LABEL_REMAP, SENTIMENT_LABELS |
| VALID_LABELS = set(SENTIMENT_LABELS.classes) |
|
|
|
|
| def _deduplicate_articles(samples: list[dict]) -> list[dict]: |
| """Remove samples with duplicate article text, keeping the first.""" |
| seen_texts = set() |
| deduped = [] |
| removed = 0 |
| for s in samples: |
| if s["text"] in seen_texts: |
| removed += 1 |
| continue |
| seen_texts.add(s["text"]) |
| deduped.append(s) |
| print(f"Deduplicated articles: removed {removed}, kept {len(deduped)}") |
| return deduped |
|
|
|
|
| def _remap_labels(samples: list[dict]) -> list[dict]: |
| """Remap non-standard labels and drop entities with unmappable labels.""" |
| mapping = LABEL_REMAP.mapping |
| remapped_count = 0 |
| dropped_count = 0 |
|
|
| for s in samples: |
| cleaned_entities = [] |
| for e in s["entities"]: |
| label = e["label"] |
| if label in VALID_LABELS: |
| cleaned_entities.append(e) |
| elif label in mapping: |
| e["label"] = mapping[label] |
| cleaned_entities.append(e) |
| remapped_count += 1 |
| else: |
| dropped_count += 1 |
| s["entities"] = cleaned_entities |
|
|
| print(f"Labels remapped: {remapped_count}, dropped (unmappable): {dropped_count}") |
| return samples |
|
|
|
|
| def _fix_position_text(samples: list[dict]) -> list[dict]: |
| """Overwrite position_text with actual span from article text when case differs.""" |
| fixed = 0 |
| for s in samples: |
| text = s["text"] |
| for e in s["entities"]: |
| for p in e["positions"]: |
| end = p["offset"] + p["length"] |
| if end > len(text): |
| continue |
| actual = text[p["offset"]:end] |
| if actual != p["position_text"] and actual.lower() == p["position_text"].lower(): |
| p["position_text"] = actual |
| fixed += 1 |
| print(f"Position text case mismatches fixed: {fixed}") |
| return samples |
|
|
|
|
| def _merge_entities(samples: list[dict]) -> list[dict]: |
| """Merge entity records sharing (entity_text, label) within a sample. |
| |
| Collects all positions into a single entity record per unique |
| (entity_text.lower(), label) pair. |
| """ |
| total_merged = 0 |
| different_label = 0 |
|
|
| for s in samples: |
| merged: dict[str, dict] = {} |
| label_seen: dict[str, str] = {} |
|
|
| for e in s["entities"]: |
| key = e["entity_text"].lower() |
|
|
| if key in label_seen and label_seen[key] != e["label"]: |
| different_label += 1 |
|
|
| merge_key = (key, e["label"]) |
| if merge_key in merged: |
| merged[merge_key]["positions"].extend(e["positions"]) |
| total_merged += 1 |
| else: |
| merged[merge_key] = { |
| "entity_id": e["entity_id"], |
| "entity_text": e["entity_text"], |
| "entity_type": e["entity_type"], |
| "positions": list(e["positions"]), |
| "label": e["label"], |
| } |
| label_seen[key] = e["label"] |
|
|
| s["entities"] = list(merged.values()) |
|
|
| print(f"Entities merged: {total_merged}, different-label pairs: {different_label}") |
| return samples |
|
|
|
|
| def _deduplicate_positions(samples: list[dict]) -> list[dict]: |
| """Remove exact-duplicate positions (same offset + length) within each entity.""" |
| removed = 0 |
| for s in samples: |
| for e in s["entities"]: |
| seen = set() |
| unique = [] |
| for p in e["positions"]: |
| span = (p["offset"], p["length"]) |
| if span in seen: |
| removed += 1 |
| continue |
| seen.add(span) |
| unique.append(p) |
| e["positions"] = unique |
| print(f"Exact-duplicate positions removed: {removed}") |
| return samples |
|
|
|
|
| def _resolve_same_offset(samples: list[dict]) -> list[dict]: |
| """For positions sharing the same offset, keep the longest span.""" |
| resolved = 0 |
| for s in samples: |
| for e in s["entities"]: |
| by_offset: dict[int, dict] = {} |
| for p in e["positions"]: |
| off = p["offset"] |
| if off not in by_offset or p["length"] > by_offset[off]["length"]: |
| if off in by_offset: |
| resolved += 1 |
| by_offset[off] = p |
| e["positions"] = sorted(by_offset.values(), key=lambda p: p["offset"]) |
| print(f"Same-offset positions resolved (kept longest): {resolved}") |
| return samples |
|
|
|
|
| def _resolve_partial_overlaps(samples: list[dict]) -> list[dict]: |
| """Resolve positions that partially overlap (different offset, shared characters). |
| |
| Positions are already sorted by offset. When two positions overlap, |
| the longer span is kept and the shorter one is discarded. |
| """ |
| resolved = 0 |
| for s in samples: |
| for e in s["entities"]: |
| positions = sorted(e["positions"], key=lambda p: p["offset"]) |
| kept: list[dict] = [] |
| for p in positions: |
| if not kept: |
| kept.append(p) |
| continue |
| prev = kept[-1] |
| prev_end = prev["offset"] + prev["length"] |
| if p["offset"] < prev_end: |
| resolved += 1 |
| if p["length"] > prev["length"]: |
| kept[-1] = p |
| else: |
| kept.append(p) |
| e["positions"] = kept |
| print(f"Partial overlaps resolved (kept longest): {resolved}") |
| return samples |
|
|
|
|
| def preprocess(samples: list[dict]) -> list[dict]: |
| """Run the full preprocessing pipeline on raw samples.""" |
| samples = _deduplicate_articles(samples) |
| samples = _remap_labels(samples) |
| samples = _fix_position_text(samples) |
| samples = _merge_entities(samples) |
| samples = _deduplicate_positions(samples) |
| samples = _resolve_same_offset(samples) |
| samples = _resolve_partial_overlaps(samples) |
|
|
| total_entities = sum(len(s["entities"]) for s in samples) |
| print(f"Preprocessing complete: {len(samples)} samples, {total_entities} entities") |
| return samples |
|
|
|
|
| def save_jsonl(samples: list[dict], path: str | Path) -> None: |
| """Write preprocessed samples as JSONL (one JSON object per line).""" |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with open(path, "w", encoding="utf-8") as f: |
| for s in samples: |
| f.write(json.dumps(s, ensure_ascii=False) + "\n") |
| print(f"Saved {len(samples)} samples to {path}") |
|
|
|
|
| def main( |
| input_path: str | Path = "data/data_raw.json", |
| output_path: str | Path = "data/data_preprocessed.jsonl", |
| ) -> list[dict]: |
| """Load raw data, preprocess, and save.""" |
| with open(input_path, "r", encoding="utf-8") as f: |
| raw = json.load(f) |
|
|
| print(f"Loaded {len(raw)} raw samples from {input_path}") |
| samples = preprocess(raw) |
| save_jsonl(samples, output_path) |
| return samples |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
