mirror of
https://github.com/oobabooga/text-generation-webui.git
synced 2026-04-06 07:03:37 +00:00
39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
from pathlib import Path
|
|
|
|
from modules import shared, utils
|
|
from modules.utils import sanitize_filename
|
|
from modules.text_generation import get_encoded_length
|
|
|
|
|
|
def load_prompt(fname):
|
|
if not fname:
|
|
# Create new file
|
|
new_name = utils.current_time()
|
|
prompt_path = shared.user_data_dir / "logs" / "notebook" / f"{new_name}.txt"
|
|
prompt_path.parent.mkdir(parents=True, exist_ok=True)
|
|
initial_content = "In this story,"
|
|
prompt_path.write_text(initial_content, encoding='utf-8')
|
|
|
|
# Update settings to point to new file
|
|
shared.settings['prompt-notebook'] = new_name
|
|
|
|
return initial_content
|
|
|
|
fname = sanitize_filename(fname)
|
|
file_path = shared.user_data_dir / 'logs' / 'notebook' / f'{fname}.txt'
|
|
if file_path.exists():
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
text = f.read()
|
|
text = text.rstrip()
|
|
|
|
return text
|
|
else:
|
|
return ''
|
|
|
|
|
|
def count_tokens(text):
|
|
try:
|
|
tokens = get_encoded_length(text)
|
|
return str(tokens)
|
|
except Exception:
|
|
return '0'
|