kaizenkodo/ChatMastermind
2
0
Fork 0
Code Issues 2 Pull Requests Packages Projects Releases Wiki Activity

Optimize tags retrival in TXT files.

Browse Source
This commit is contained in:
Oleksandr Kozachuk 2023-05-07 14:43:56 +02:00
parent c5fd466dda
commit 63a202376d
1 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
chatmastermind/storage.py
View File

@ -5,9 +5,11 @@ from .utils import terminal_width, append_message, message_to_chat
from typing import List, Dict, Any, Optional from typing import List, Dict, Any, Optional
def read_file(fname: str) -> Dict[str, Any]: def read_file(fname: str, tags_only: bool = False) -> Dict[str, Any]:
with open(fname, "r") as fd: with open(fname, "r") as fd:
text = fd.read().split('\n') if tags_only:
return {"tags": [x.strip() for x in fd.readline().strip().split(':')[1].strip().split(',')]}
text = fd.read().strip().split('\n')
tags = [x.strip() for x in text.pop(0).split(':')[1].strip().split(',')] tags = [x.strip() for x in text.pop(0).split(':')[1].strip().split(',')]
question_idx = text.index("=== QUESTION ===") + 1 question_idx = text.index("=== QUESTION ===") + 1
answer_idx = text.index("==== ANSWER ====") answer_idx = text.index("==== ANSWER ====")
@ -91,7 +93,7 @@ def get_tags(config: Dict[str, Any], prefix: Optional[str]) -> List[str]:
with open(file, 'r') as f: with open(file, 'r') as f:
data = yaml.load(f, Loader=yaml.FullLoader) data = yaml.load(f, Loader=yaml.FullLoader)
elif file.suffix == '.txt': elif file.suffix == '.txt':
data = read_file(file) data = read_file(file, tags_only=True)
else: else:
continue continue
for tag in data.get('tags', []): for tag in data.get('tags', []):