save processed documents and model loading

same as https://github.com/ArneBinder/pie-document-level/pull/213/commits/34327b71b6b1a50341d003311888402f6705b3bc

app.py

@@ -1,7 +1,7 @@
 import json
 import logging
 from functools import partial
-from typing import Optional, Tuple
+from typing import Any, Optional, Tuple
 
 import gradio as gr
 from pie_modules.document.processing import tokenize_document
@@ -22,6 +22,13 @@ logger = logging.getLogger(__name__)
 RENDER_WITH_DISPLACY = "displaCy + highlighted arguments"
 RENDER_WITH_PRETTY_TABLE = "Pretty Table"
 
+DEFAULT_MODEL_NAME = "ArneBinder/sam-pointer-bart-base-v0.3"
+DEFAULT_MODEL_REVISION = "76300f8e534e2fcf695f00cb49bba166739b8d8a"
+# local path
+# DEFAULT_MODEL_NAME = "models/dataset-sciarg/task-ner_re/v0.3/2024-05-28_23-33-46"
+# DEFAULT_MODEL_REVISION = None
+DEFAULT_EMBEDDING_MODEL_NAME = "allenai/scibert_scivocab_uncased"
+
 
 def embed_text_annotations(
     document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
@@ -99,6 +106,10 @@ def predict(
         # convert keys to str because JSON keys must be strings
         adu_embeddings_dict = {str(k._id): v.detach().tolist() for k, v in adu_embeddings.items()}
         document_dict["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."
+        )
 
     # Return as dict and JSON string. The latter is required because the JSON component converts floats
     # to ints which destroys de-serialization of the document (the scores of the annotations need to be floats)
@@ -123,6 +134,29 @@ def render(document_txt: str, render_with: str, render_kwargs_json: str) -> str:
     return html
 
 
+def add_to_index(
+    output_txt: str, doc_id: str, processed_documents: dict, vector_store: Any
+) -> None:
+    try:
+        if doc_id in processed_documents:
+            gr.Warning(f"Document {doc_id} already in index. Overwriting.")
+        output = json.loads(output_txt)
+        # get the embeddings from the output and remove them from the output
+        embeddings = output.pop("embeddings")
+        # save the processed document to the index
+        processed_documents[doc_id] = output
+        # save the embeddings to the vector store
+        for adu_id, embedding in embeddings.items():
+            emb_id = f"{doc_id}:{adu_id}"
+            # TODO: save embedding to vector store at emb_id (embedding is a list of 768 floats)
+
+        gr.Info(
+            f"Added document {doc_id} to index (index contains {len(processed_documents)} entries). (NOT YET IMPLEMENTED)"
+        )
+    except Exception as e:
+        raise gr.Error(f"Failed to add document {doc_id} to index: {e}")
+
+
 def open_accordion():
     return gr.Accordion(open=True)
 
@@ -131,30 +165,52 @@ def close_accordion():
     return gr.Accordion(open=False)
 
 
-def main():
+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
 
-    model_name_or_path = "ArneBinder/sam-pointer-bart-base-v0.3"
-    # W&B run: https://wandb.ai/arne/dataset-sciarg-task-ner_re-v0.3-training/runs/prik91di
-    revision = "76300f8e534e2fcf695f00cb49bba166739b8d8a"
-    # local path
-    # model_name_or_path = "models/dataset-sciarg/task-ner_re/v0.3/2024-05-28_23-33-46"
-    # revision = None
+
+def main():
 
     example_text = "Scholarly Argumentation Mining (SAM) has recently gained attention due to its potential to help scholars with the rapid growth of published scientific literature. It comprises two subtasks: argumentative discourse unit recognition (ADUR) and argumentative relation extraction (ARE), both of which are challenging since they require e.g. the integration of domain knowledge, the detection of implicit statements, and the disambiguation of argument structure. While previous work focused on dataset construction and baseline methods for specific document sections, such as abstract or results, full-text scholarly argumentation mining has seen little progress. In this work, we introduce a sequential pipeline model combining ADUR and ARE for full-text SAM, and provide a first analysis of the performance of pretrained language models (PLMs) on both subtasks. We establish a new SotA for ADUR on the Sci-Arg corpus, outperforming the previous best reported result by a large margin (+7% F1). We also present the first results for ARE, and thus for the full AM pipeline, on this benchmark dataset. Our detailed error analysis reveals that non-contiguous ADUs as well as the interpretation of discourse connectors pose major challenges and that data annotation needs to be more consistent."
 
-    print("loading argumentation mining model ...")
-    pipeline = AutoPipeline.from_pretrained(
-        model_name_or_path,
-        device=-1,
-        num_workers=0,
-        taskmodule_kwargs=dict(revision=revision),
-        model_kwargs=dict(revision=revision),
+    print("Loading argumentation model ...")
+    argumentation_model = load_argumentation_model(
+        model_name=DEFAULT_MODEL_NAME, revision=DEFAULT_MODEL_REVISION
     )
 
-    print("loading SciBERT embedding model ...")
-    embedding_model = AutoModel.from_pretrained("allenai/scibert_scivocab_uncased")
-    embedding_tokenizer = AutoTokenizer.from_pretrained("allenai/scibert_scivocab_uncased")
-
     default_render_kwargs = {
         "entity_options": {
             # we need to convert the keys to uppercase because the spacy rendering function expects them in uppercase
@@ -180,18 +236,49 @@ def main():
         },
     }
 
+    # TODO: setup the vector store
+    vector_store = None
+
     with gr.Blocks() as demo:
+        processed_documents_state = gr.State(dict())
+        vector_store_state = gr.State(vector_store)
+        # wrap the pipeline and the embedding model/tokenizer in a tuple to avoid that it gets called
+        models_state = gr.State((argumentation_model, None, None))
         with gr.Row():
             with gr.Column(scale=1):
+                doc_id = gr.Textbox(
+                    label="Document ID",
+                    value="user_input",
+                )
                 text = gr.Textbox(
-                    label="Input Text",
+                    label="Text",
                     lines=20,
                     value=example_text,
                 )
+                with gr.Accordion("Model Configuration", open=False):
+                    model_name = gr.Textbox(
+                        label="Model Name",
+                        value=DEFAULT_MODEL_NAME,
+                    )
+                    model_revision = gr.Textbox(
+                        label="Model Revision",
+                        value=DEFAULT_MODEL_REVISION,
+                    )
+                    embedding_model_name = gr.Textbox(
+                        label=f"Embedding Model Name (e.g. {DEFAULT_EMBEDDING_MODEL_NAME})",
+                        value="",
+                    )
+                    load_models_btn = gr.Button("Load Models")
+                    load_models_btn.click(
+                        fn=load_models,
+                        inputs=[model_name, model_revision, embedding_model_name],
+                        outputs=models_state,
+                    )
 
-                predict_btn = gr.Button("Predict")
+                predict_btn = gr.Button("Analyse")
 
                 output_txt = gr.Textbox(visible=False)
+                add_to_index_btn = gr.Button("Add current result to Index")
 
             with gr.Column(scale=1):
 
@@ -216,14 +303,16 @@ def main():
         render_button_kwargs = dict(
             fn=render, inputs=[output_txt, render_as, render_kwargs], outputs=rendered_output
         )
+
+        def _predict(
+            text: str,
+            models: Tuple[Pipeline, Optional[PreTrainedModel], Optional[PreTrainedTokenizer]],
+        ) -> Tuple[dict, str]:
+            return predict(text, *models)
+
         predict_btn.click(open_accordion, inputs=[], outputs=[output_accordion]).then(
-            fn=partial(
-                predict,
-                pipeline=pipeline,
-                embedding_model=embedding_model,
-                embedding_tokenizer=embedding_tokenizer,
-            ),
-            inputs=text,
+            fn=_predict,
+            inputs=[text, models_state],
             outputs=[output_json, output_txt],
             api_name="predict",
         ).success(**render_button_kwargs).success(
@@ -231,6 +320,12 @@ def main():
         )
         render_btn.click(**render_button_kwargs, api_name="render")
 
+        add_to_index_btn.click(
+            fn=add_to_index,
+            inputs=[output_txt, doc_id, processed_documents_state, vector_store_state],
+            outputs=[],
+        )
+
         js = """
         () => {
             function maybeSetColor(entity, colorAttributeKey, colorDictKey) {