Files
fn_registry/python/functions/datascience/mrebel_load_model.py
T
egutierrez faac610745 feat: extraccion masiva footprint_aurgi (41 funcs + 4 types + stack Docker geo)
Extrae al registry funciones del proyecto interno footprint_aurgi:
- core (6): slugify_ascii, normalize_for_join, cp_provincia_es, infer_provincia_from_cp, safe_read_csv_fallback, csv_to_parquet_duckdb
- geo puras (7): haversine_km, point_in_ring, point_in_polygon, point_in_polygons_bbox, polygon_bbox, extent_with_padding, distance_bucket
- geo I/O (4): load_geojson_polygons, load_boundary_gdf, add_basemap_osm, add_basemap_with_timeout
- valhalla client (4): valhalla_route, valhalla_isochrone, valhalla_isochrones_async, valhalla_matrix_1_to_n
- datascience stats (7): trimmed_mean, geometric_mean, detect_distribution_type, best_central_tendency, summary_stats, kde_density_levels, alpha_shape_concave_hull
- datascience fuzzy (3): fuzzy_merge_adaptive (rapidfuzz), words_to_dataset, remove_words_from_column
- datascience viz (2): plot_kde_2d, plot_heatmap_log
- infra (4): compress_pdf_ghostscript, render_table_page_pdfpages, add_header_logo, osm2pgsql_ingest
- pipelines (4): setup_geo_stack_docker, compute_centers_reachability, generate_isochrones_by_zone, count_points_per_zone
- types geo (4): LonLat, BBox, IsochroneRequest, Centro

Incluye:
- apps/footprint_geo_stack/ (PostGIS + Martin + Valhalla via docker-compose)
- 131/132 tests pasan (1 skip esperado: osm2pgsql en PATH)
- Issue tracker dev/issues/0052-footprint-aurgi-extraction.md
- Atribucion uniforme: source_repo internal:footprint_aurgi, source_license internal-aurgi
- Build con 9 agentes en paralelo (8 wave 1 + 1 wave 2 pipelines)

Tambien commitea trabajo previo no commiteado: aggregate_extraction_results, chunk_with_overlap, clean_pdf_text, merge_entity_aliases, extract_graph_gliner2, extract_relations_mrebel, extract_triples_spacy_es, gliner2/mrebel/marianmt/rebel/spacy_es load_model, parse_rebel_output, translate_es_to_en, issue 0050/0051.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-04 23:35:22 +02:00

70 lines
2.5 KiB
Python

"""Carga (y cachea) el modelo mREBEL para extraccion de relaciones multilingue."""
from __future__ import annotations
from typing import Any
# Cache global: (model_name, src_lang) -> (tokenizer, model)
_MODEL_CACHE: dict[tuple[str, str], tuple[Any, Any]] = {}
def mrebel_load_model(
model_name: str = "Babelscape/mrebel-large",
src_lang: str = "es_XX",
tgt_lang: str = "tp_XX",
) -> tuple[Any, Any]:
"""Loads (and caches) the mREBEL tokenizer and model.
mREBEL is a multilingual seq2seq model (mBART-based, ~600M params, ~2.4 GB)
for relation extraction. It supports 30+ languages via language codes
(``src_lang``).
LICENSE NOTICE: Babelscape/mrebel-large is licensed under CC BY-NC-SA 4.0
(Creative Commons Non-Commercial Share-Alike). Do NOT use in commercial
products without replacing this model with a commercially-licensed
alternative (e.g. Babelscape/rebel-large which is Apache 2.0 but
English-only).
The first call downloads the model from HuggingFace Hub (~2.4 GB).
Subsequent calls with the same ``(model_name, src_lang)`` return the
cached instance without re-loading.
Args:
model_name: HuggingFace Hub model ID. Default is the large variant.
src_lang: Source language code for the mBART tokenizer, e.g.
``"es_XX"`` (Spanish), ``"en_XX"`` (English), ``"fr_XX"`` (French).
tgt_lang: Target language token for the decoder (always ``"tp_XX"``
for the triplet format — only change if using a custom checkpoint).
Returns:
Tuple ``(tokenizer, model)`` both ready for inference with
``model.generate(...)`` and ``tokenizer.decode(...)``.
Raises:
ImportError: if ``transformers`` is not installed.
OSError: if the model cannot be downloaded or loaded from disk.
"""
cache_key = (model_name, src_lang)
cached = _MODEL_CACHE.get(cache_key)
if cached is not None:
return cached
try:
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
except ImportError as exc:
raise ImportError(
"transformers no esta instalado. Instalalo con "
"`uv pip install transformers` o `uv pip install -e '.[nlp]'`."
) from exc
tokenizer = AutoTokenizer.from_pretrained(
model_name,
src_lang=src_lang,
tgt_lang=tgt_lang,
)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
model.eval()
_MODEL_CACHE[cache_key] = (tokenizer, model)
return tokenizer, model