text-generator-v2 / datamanager.py
Cat0125
add train tab, improve quality
8e637c7
raw
history blame contribute delete
No virus
5.26 kB
import json
import pickle
import time
from files import read_file, read_lines
models = json.load(open("models/models.json"))
TEXT_PATH = 'models/%s/text.txt'
FILENAME_V1 = 'models/%s/data.pkl'
FILENAME_V2 = 'models/%s/data2.pkl'
FILENAME_V3 = 'models/%s/data3.pkl'
def get_texts(model_name):
"""
This function returns the lines of text associated with a given model name.
:param model_name: The name of a model that has been defined in the `models` dictionary. This
function is designed to retrieve the texts associated with a particular model
:return: The function `get_texts` is returning the text data from a specific model, which is
identified by its name. The text data is obtained by calling the `read_lines` function on the `text`
attribute of the specified model.
"""
return read_lines(TEXT_PATH % model_name)
def get_text(model_name):
return read_file(TEXT_PATH % model_name)
def set_data(model_name, data):
"""
This function saves data to a file using the pickle module, with the filename specified by the
model_name argument.
:param model_name: The name of the model for which the data is being set
:param data: The data that needs to be saved for the given model. It could be any Python object such
as a list, dictionary, or a trained model
"""
print(f'Writing data for {model_name}...', end=' ')
pickle.dump(data, open(FILENAME_V1 % model_name, 'wb+'))
print('done')
def get_data(model_name):
"""
The function retrieves data from a database or a file using a model name as input.
:param model_name: The name of the model for which we want to retrieve the data
:return: The function `get_data` returns the database object for the specified `model_name`. If the
database object is already loaded in memory, it returns the cached object. Otherwise, it loads the
object from a file using `pickle.load()` and caches it for future use.
"""
if models[model_name]["db"]:
return models[model_name]["db"]
print(f'Loading model {model_name}...', end=' ')
start_time = time.time()
db = pickle.load(open(FILENAME_V1 % model_name, 'rb'))
models[model_name]["db"] = db
print("done (%ss)" % (time.time() - start_time))
return db
def set_data_v2(model_name, data):
"""
This function saves data to a file using the pickle module, with the filename specified in a
dictionary associated with the given model name.
:param model_name: The name of the model for which the data is being set
:param data: The data that needs to be saved to a file using the pickle module
"""
print(f'Writing data for {model_name}...', end=' ')
pickle.dump(data, open(FILENAME_V2 % model_name, 'wb+'))
print('done')
def get_data_v2(model_name):
"""
This function returns a database object for a given model name, either by loading it from a file or
returning a cached version.
:param model_name: The name of the model for which we want to retrieve the data
:return: a database object for the given model name. If the database object is already loaded in the
models dictionary, it returns the object from the dictionary. Otherwise, it loads the object from a
pickle file and stores it in the dictionary before returning it.
"""
if models[model_name]["db2"]:
return models[model_name]["db2"]
print(f'Loading model {model_name}...', end=' ')
start_time = time.time()
db = pickle.load(open(FILENAME_V2 % model_name, 'rb'))
models[model_name]["db2"] = db
print("done (%ss)" % (time.time() - start_time))
return db
def set_data_v3(model_name, data):
"""
This function saves data to a file using the pickle module, with the filename specified by the
model_name argument.
:param model_name: The name of the model for which the data is being set
:param data: The data parameter is the data that needs to be saved to a file using the pickle
module. The data can be of any type, such as a list, dictionary, or object. The function saves the
data to a file specified by the model_name parameter. The filename is obtained from the models
dictionary
"""
print(f'Writing data for {model_name}...', end=' ')
pickle.dump(data, open(FILENAME_V3 % model_name, 'wb+'))
print('done')
def get_data_v3(model_name):
"""
This function loads a database file for a given model and returns it, while also caching it for
future use.
:param model_name: a string representing the name of a model
:return: The function `get_data_v3` returns the database object for the given `model_name`. If the
database object is already loaded in the `models` dictionary, it returns the cached object.
Otherwise, it loads the object from the file specified in the `models` dictionary, caches it in the
`models` dictionary, and returns it.
"""
if models[model_name]["db3"]:
return models[model_name]["db3"]
print(f'Loading model {model_name}...', end=' ')
start_time = time.time()
db = pickle.load(open(FILENAME_V3 % model_name, 'rb'))
models[model_name]["db3"] = db
print("done (%ss)" % (time.time() - start_time))
return db