Upload files to "OHS"

OHS/Script.py

@@ -0,0 +1,77 @@
+import json
+import numpy as np
+from typing import List, Dict, Tuple, Any, Optional
+from dataclasses import dataclass
+from pathlib import Path
+
+@dataclass
+class ObjectionHandleScript:
+    def __init__(self, id: int, script: str, target_problem: str, target_customer: str):
+
+        self.id: int = id
+        self.script: str = script
+        self.target_problem: str = target_problem
+        self.target_customer: str = target_customer
+
+    def get_id(self):
+        return self.id
+
+    def get_script(self):
+        return self.script
+
+    def get_target_problem(self):
+        return self.target_problem
+
+    def get_target_customer(self):
+        """
+        用于Reranker阶段的重排序
+        """
+        return self.target_customer
+
+    def get_document_text(self) -> str:
+        """
+        用于Embedding阶段召回和Reranker阶段的重排序
+        """
+        return f"目标解决问题: {self.target_problem}\n目标客户特征: {self.target_customer}"
+
+class ScriptDataset:
+    def __init__(self, json_file: str):
+        self.json_file = Path(json_file)
+        
+        with open(json_file, 'r', encoding='utf-8') as f:
+            raw_data = json.load(f)
+        assert isinstance(raw_data, list), "JSON file should contain a list of objects"
+
+        self.scripts = []
+        for idx, e in enumerate(raw_data):
+            self.scripts.append(ObjectionHandleScript(
+                id=idx,
+                script=e['话术本身'],
+                target_problem=e['目标解决问题'],
+                target_customer=e['目标客户特征']
+            ))
+
+        self.index_file = self.json_file.with_suffix('.faiss')
+        self.vectors_file = self.json_file.with_suffix('.vectors.npy')
+        self.metadata_file = self.json_file.with_suffix('.metadata.pkl')
+        self.faiss_index = None
+        self.document_vectors = None
+
+    def get_script(self, idx) -> ObjectionHandleScript:
+        return self.scripts[idx]
+
+    def __len__(self):
+        return len(self.scripts)
+
+    def __getitem__(self, idx):
+        return self.scripts[idx]
+
+    def get_all_document_texts(self) -> List[str]:
+        return [script.get_document_text() for script in self.scripts]
+
+    def get_all_script_texts(self) -> List[str]:
+        return [script.get_script() for script in self.scripts]
+
+if __name__=="__main__":
+    json_path = "scripts_deduplicated.json"
+    dataset = ScriptDataset(json_path)