|
import logging |
|
from typing import Dict, List, Optional, Tuple |
|
|
|
import gradio as gr |
|
from annotation_utils import labeled_span_to_id |
|
from pie_modules.document.processing import tokenize_document |
|
from pie_modules.documents import TokenDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions |
|
from pytorch_ie import Pipeline |
|
from pytorch_ie.annotations import LabeledSpan |
|
from pytorch_ie.auto import AutoPipeline |
|
from pytorch_ie.documents import TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions |
|
from transformers import AutoModel, AutoTokenizer, PreTrainedModel, PreTrainedTokenizer |
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
def _embed_text_annotations( |
|
document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions, |
|
model: PreTrainedModel, |
|
tokenizer: PreTrainedTokenizer, |
|
text_layer_name: str, |
|
) -> Dict[LabeledSpan, List[float]]: |
|
|
|
document = document.copy() |
|
|
|
document[text_layer_name].extend(document[text_layer_name].predictions.clear()) |
|
added_annotations = [] |
|
tokenizer_kwargs = { |
|
"max_length": 512, |
|
"stride": 64, |
|
"truncation": True, |
|
"return_overflowing_tokens": True, |
|
} |
|
tokenized_documents = tokenize_document( |
|
document, |
|
tokenizer=tokenizer, |
|
result_document_type=TokenDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions, |
|
partition_layer="labeled_partitions", |
|
added_annotations=added_annotations, |
|
strict_span_conversion=False, |
|
**tokenizer_kwargs, |
|
) |
|
|
|
model_inputs = tokenizer(document.text, return_tensors="pt", **tokenizer_kwargs) |
|
|
|
model_inputs.pop("overflow_to_sample_mapping", None) |
|
assert len(model_inputs.encodings) == len(tokenized_documents) |
|
model_output = model(**model_inputs) |
|
|
|
|
|
embeddings = {} |
|
for batch_idx in range(len(model_output.last_hidden_state)): |
|
text2tok_ann = added_annotations[batch_idx][text_layer_name] |
|
tok2text_ann = {v: k for k, v in text2tok_ann.items()} |
|
for tok_ann in tokenized_documents[batch_idx].labeled_spans: |
|
|
|
if tok_ann.start == tok_ann.end: |
|
continue |
|
|
|
embedding = model_output.last_hidden_state[batch_idx, tok_ann.start : tok_ann.end].max( |
|
dim=0 |
|
)[0] |
|
text_ann = tok2text_ann[tok_ann] |
|
|
|
if text_ann in embeddings: |
|
logger.warning( |
|
f"Overwriting embedding for annotation '{text_ann}' (do you use striding?)" |
|
) |
|
embeddings[text_ann] = embedding |
|
|
|
return embeddings |
|
|
|
|
|
def _annotate( |
|
document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions, |
|
pipeline: Pipeline, |
|
embedding_model: Optional[PreTrainedModel] = None, |
|
embedding_tokenizer: Optional[PreTrainedTokenizer] = None, |
|
) -> None: |
|
|
|
|
|
pipeline(document) |
|
|
|
if embedding_model is not None and embedding_tokenizer is not None: |
|
adu_embeddings = _embed_text_annotations( |
|
document=document, |
|
model=embedding_model, |
|
tokenizer=embedding_tokenizer, |
|
text_layer_name="labeled_spans", |
|
) |
|
|
|
adu_embeddings_dict = { |
|
labeled_span_to_id(k): v.detach().tolist() for k, v in adu_embeddings.items() |
|
} |
|
document.metadata["embeddings"] = adu_embeddings_dict |
|
else: |
|
gr.Warning( |
|
"No embedding model provided. Skipping embedding extraction. You can load an embedding " |
|
"model in the 'Model Configuration' section." |
|
) |
|
|
|
|
|
def create_and_annotate_document( |
|
text: str, |
|
doc_id: str, |
|
models: Tuple[Pipeline, Optional[PreTrainedModel], Optional[PreTrainedTokenizer]], |
|
) -> TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions: |
|
"""Create a TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions from the provided |
|
text, annotate it, and add it to the index. |
|
|
|
Parameters: |
|
text: The text to process. |
|
doc_id: The ID of the document. |
|
models: A tuple containing the prediction pipeline and the embedding model and tokenizer. |
|
|
|
Returns: |
|
The processed document. |
|
""" |
|
|
|
try: |
|
document = TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions( |
|
id=doc_id, text=text, metadata={} |
|
) |
|
|
|
document.labeled_partitions.append(LabeledSpan(start=0, end=len(text), label="text")) |
|
|
|
_annotate( |
|
document=document, |
|
pipeline=models[0], |
|
embedding_model=models[1], |
|
embedding_tokenizer=models[2], |
|
) |
|
|
|
return document |
|
except Exception as e: |
|
raise gr.Error(f"Failed to process text: {e}") |
|
|
|
|
|
def load_argumentation_model(model_name: str, revision: Optional[str] = None) -> Pipeline: |
|
try: |
|
model = AutoPipeline.from_pretrained( |
|
model_name, |
|
device=-1, |
|
num_workers=0, |
|
taskmodule_kwargs=dict(revision=revision), |
|
model_kwargs=dict(revision=revision), |
|
) |
|
except Exception as e: |
|
raise gr.Error(f"Failed to load argumentation model: {e}") |
|
gr.Info(f"Loaded argumentation model: model_name={model_name}, revision={revision})") |
|
return model |
|
|
|
|
|
def load_embedding_model(model_name: str) -> Tuple[PreTrainedModel, PreTrainedTokenizer]: |
|
try: |
|
embedding_model = AutoModel.from_pretrained(model_name) |
|
embedding_tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
except Exception as e: |
|
raise gr.Error(f"Failed to load embedding model: {e}") |
|
gr.Info(f"Loaded embedding model: model_name={model_name})") |
|
return embedding_model, embedding_tokenizer |
|
|
|
|
|
def load_models( |
|
model_name: str, revision: Optional[str] = None, embedding_model_name: Optional[str] = None |
|
) -> Tuple[Pipeline, Optional[PreTrainedModel], Optional[PreTrainedTokenizer]]: |
|
argumentation_model = load_argumentation_model(model_name, revision) |
|
embedding_model = None |
|
embedding_tokenizer = None |
|
if embedding_model_name is not None and embedding_model_name.strip(): |
|
embedding_model, embedding_tokenizer = load_embedding_model(embedding_model_name) |
|
|
|
return argumentation_model, embedding_model, embedding_tokenizer |
|
|