tcsenpai/pensieve
1
0
Fork 0
mirror of https://github.com/tcsenpai/pensieve.git synced 2025-06-10 04:57:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: yaml parse related

Browse Source
This commit is contained in:
arkohut 2024-09-09 23:39:04 +08:00
parent 3912d165f6
commit a30fe62bc3
2 changed files with 3 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
memos/config.py
View File

@ -102,15 +102,13 @@ yaml.add_representer(OrderedDict, dict_representer)
def secret_str_representer(dumper, data): def secret_str_representer(dumper, data):
return dumper.represent_scalar("tag:yaml.org,2002:str", data.get_secret_value()) return dumper.represent_scalar("tag:yaml.org,2002:str", data.get_secret_value())
# Custom constructor for SecretStr # Custom constructor for SecretStr
def secret_str_constructor(loader, node): def secret_str_constructor(loader, node):
value = loader.construct_scalar(node) value = loader.construct_scalar(node)
return SecretStr(value) return SecretStr(value)
# Register the representer and constructor only for specific fields
yaml.add_representer(SecretStr, secret_str_representer) yaml.add_representer(SecretStr, secret_str_representer)
yaml.add_constructor("tag:yaml.org,2002:str", secret_str_constructor)
def create_default_config(): def create_default_config():

4
memos/plugins/ocr/main.py
View File

@ -183,10 +183,10 @@ def init_plugin(config):
ocr_config['Cls']['model_path'] = os.path.join(model_dir, os.path.basename(ocr_config['Cls']['model_path'])) ocr_config['Cls']['model_path'] = os.path.join(model_dir, os.path.basename(ocr_config['Cls']['model_path']))
ocr_config['Rec']['model_path'] = os.path.join(model_dir, os.path.basename(ocr_config['Rec']['model_path'])) ocr_config['Rec']['model_path'] = os.path.join(model_dir, os.path.basename(ocr_config['Rec']['model_path']))
# Save the updated config to a temporary file # Save the updated config to a temporary file with strings wrapped in double quotes
temp_config_path = os.path.join(os.path.dirname(__file__), "temp_ppocr.yaml") temp_config_path = os.path.join(os.path.dirname(__file__), "temp_ppocr.yaml")
with open(temp_config_path, 'w') as f: with open(temp_config_path, 'w') as f:
yaml.safe_dump(ocr_config, f) yaml.safe_dump(ocr_config, f, default_style='"')
ocr = RapidOCR(config_path=temp_config_path) ocr = RapidOCR(config_path=temp_config_path)
thread_pool = ThreadPoolExecutor(max_workers=concurrency) thread_pool = ThreadPoolExecutor(max_workers=concurrency)