feat: integrate StyleRewriting

app.py

@@ -6,6 +6,8 @@ import os
 import penman
 import sys
 import time
+import torch
+from transformers import pipeline
 
 if not os.path.exists("amr-tst-indo"):
     Repo.clone_from("https://github.com/AbdiHaryadi/amr-tst-indo.git", "amr-tst-indo")
@@ -13,6 +15,7 @@ sys.path.append("./amr-tst-indo")
 
 from text_to_amr import TextToAMR
 from style_detector import StyleDetector
+from style_rewriting import StyleRewriting
 
 amr_parsing_model_name = "mbart-en-id-smaller-indo-amr-parsing-translated-nafkhan"
 snapshot_download(
@@ -34,6 +37,25 @@ sd = StyleDetector(
     model_path="./model-best.pt"
 )
 
+device_type = "cuda" if torch.cuda.is_available() else "cpu"
+clf_pipeline = pipeline(
+    "text-classification",
+    model="abdiharyadi/roberta-base-indonesian-522M-with-sa-william-dataset",
+    device=device_type
+)
+gdown.download(
+    "https://drive.google.com/uc?id=15KctCcsHgTFMUh_tWNBNUiCyX56fq6p-",
+    "./fasttext_skipgram_indo.bin"
+)
+sr = StyleRewriting(
+    clf_pipeline=clf_pipeline,
+    fasttext_model_path="./fasttext_skipgram_indo.bin",
+    position_aware_concatenation=False,
+    reset_sense_strategy=False,
+    max_score_strategy=True,
+    maximize_style_words_expansion=False
+)
+
 def run(text, source_style):
     yield (
         "(Memproses ...)",
@@ -46,9 +68,16 @@ def run(text, source_style):
     start_time = time.time()
 
     # source_amr, *_ = t2a([text])
-    # source_amr.metadata = {}
-    # source_amr_display = penman.encode(source_amr)
-    source_amr_display = "(z0 / halo)"
+    source_amr = penman.decode("""
+    (z0 / dan
+        :op1 (z1 / bagus-01
+            :ARG1 (z2 / tempat)
+            :degree (z3 / sangat))
+        :op2 (z4 / bersih-01
+            :ARG1 z2))
+    """)
+    source_amr.metadata = {}
+    source_amr_display = penman.encode(source_amr)
     source_amr_display += f"\n\n({time.time() - start_time:.2f} s)"
     yield (
         source_amr_display,
@@ -85,7 +114,7 @@ def run(text, source_style):
         "(Menunggu ...)",
     )
 
-    target_amr = penman.decode("(z0 / dunia)")
+    target_amr = sr(text, source_amr, source_style, style_words)
     target_amr_display = penman.encode(target_amr)
     target_amr_display += f"\n\n({time.time() - start_time:.2f} s)"
     yield (