model

app.py

@@ -1,7 +1,14 @@
+from typing import Any, Dict
 import gradio as gr
+from pipeline import LongCorefPipeline
+from huggingface_hub import snapshot_download
 
-def greet(name):
-    return "Hello " + name + "!!"
+long_coref = LongCorefPipeline(snapshot_download("kwang2049/long-coref"))
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
+
+def fn(doc: str) -> Dict[str, Any]:
+    return long_coref(doc)
+
+
+demo = gr.Interface(fn=fn, inputs="text", outputs="text")
 demo.launch()

pipeline.py

@@ -0,0 +1,29 @@
+import os
+from typing import Dict, Any
+from long_coref.coref.prediction import CorefPredictor
+from long_coref.coref.utils import ArchiveContent
+from allennlp.common.params import Params
+
+CHECKPOINT = "coref-spanbert-large-2021.03.10"
+
+
+class LongCorefPipeline:
+    def __init__(self, path=""):
+        archive_content = ArchiveContent(
+            archive_dir=os.path.join(path, CHECKPOINT),
+            weight_path=os.path.join(path, CHECKPOINT, "weights.th"),
+            config=Params.from_file(os.path.join(path, CHECKPOINT, "config.json")),
+        )
+        self.predictor = CorefPredictor.from_extracted_archive(archive_content)
+
+    def __call__(self, data: str) -> Dict[str, Any]:
+        """
+         data args:
+              inputs (:obj: `str`)
+              date (:obj: `str`)
+        Return:
+              A :obj:`list` | `dict`: will be serialized and returned
+        """
+        # get inputs
+        prediction = self.predictor.resolve_paragraphs(data.split("\n\n"))
+        return prediction.to_dict()

requirements.txt

@@ -0,0 +1,5 @@
+long-coref @ git+https://github.com/kwang2049/long-coref.git@727612f441a91662565254a71c3bee72acf88797
+en-core-web-sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.3.0/en_core_web_sm-3.3.0-py3-none-any.whl
+thinc==8.2.5
+numpy==1.26.4
+spacy==3.7.6