Create relative_tester.py

relative_tester.py

@@ -0,0 +1,77 @@
+import torch
+import nltk
+from roberta_model_loader import roberta_model
+from feature_ref_loader import feature_mgt_ref, feature_hwt_ref
+from meta_train import net
+from regression_model_loader import regression_model
+from MMD import MMD_3_Sample_Test
+from utils import DEVICE, FeatureExtractor, HWT, MGT
+
+
+class RelativeTester:
+    def __init__(self):
+        print("Relative Tester init")
+        self.feature_extractor = FeatureExtractor(roberta_model, net)
+
+    def sents_split(self, text):
+        nltk.download("punkt", quiet=True)
+        nltk.download("punkt_tab", quiet=True)
+        sents = nltk.sent_tokenize(text)
+        return [sent for sent in sents if 5 < len(sent.split())]
+
+    def test(self, input_text, threshold=0.2, round=20):
+        print("Relative Tester test")
+        # Split the input text
+        sents = self.sents_split(input_text)
+        print("DEBUG: sents:", len(sents))
+        # Extract features
+        feature_for_sents = self.feature_extractor.process_sents(sents, False)
+        if len(feature_for_sents) <= 1:
+            # print("DEBUG: tooshort")
+            return "Too short to test! Please input more than 2 sentences."
+        # Cutoff the features
+        min_len = min(
+            len(feature_for_sents),
+            len(feature_hwt_ref),
+            len(feature_mgt_ref),
+        )
+        # Calculate MMD
+        h_u_list = []
+        p_value_list = []
+        t_list = []
+
+        for i in range(round):
+            feature_for_sents_sample = feature_for_sents[
+                torch.randperm(len(feature_for_sents))[:min_len]
+            ]
+            feature_hwt_ref_sample = feature_hwt_ref[
+                torch.randperm(len(feature_hwt_ref))[:min_len]
+            ]
+            feature_mgt_ref_sample = feature_mgt_ref[
+                torch.randperm(len(feature_mgt_ref))[:min_len]
+            ]
+            h_u, p_value, t, *rest = MMD_3_Sample_Test(
+                net.net(feature_for_sents_sample),
+                net.net(feature_hwt_ref_sample),
+                net.net(feature_mgt_ref_sample),
+                feature_for_sents_sample.view(feature_for_sents_sample.shape[0], -1),
+                feature_hwt_ref_sample.view(feature_hwt_ref_sample.shape[0], -1),
+                feature_mgt_ref_sample.view(feature_mgt_ref_sample.shape[0], -1),
+                net.sigma,
+                net.sigma0_u,
+                net.ep,
+                0.05,
+            )
+
+            h_u_list.append(h_u)
+            p_value_list.append(p_value)
+            t_list.append(t)
+
+        power = sum(h_u_list) / len(h_u_list)
+        print("DEBUG: power:", power)
+        print("DEBUG: power list:", h_u_list)
+        # Return the result
+        return "Human" if power <= threshold else "AI"
+
+
+relative_tester = RelativeTester()