42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
|
|
import json
|
|
import os
|
|
from typing import ClassVar
|
|
|
|
class Parts_Ai_DataStore:
|
|
TRAINING_DATA_DIRECTORY: ClassVar[str] = "docs/training_data"
|
|
|
|
training_data: ClassVar[object]
|
|
training_data_directory: ClassVar[str]
|
|
|
|
def __init__(self, training_data_directory):
|
|
self.training_data_directory = training_data_directory
|
|
|
|
@classmethod
|
|
def make_default(cls):
|
|
return cls(cls.TRAINING_DATA_DIRECTORY)
|
|
|
|
def load_training_data(self):
|
|
data = []
|
|
for root, dirs, files in os.walk(self.TRAINING_DATA_DIRECTORY):
|
|
for file in files:
|
|
file_path = os.path.join(root, file)
|
|
if file.endswith(".json"):
|
|
data.append({
|
|
"type": "json"
|
|
, "content": self.load_json_data(file_path)
|
|
})
|
|
elif file.endswith(".txt"):
|
|
data.append({
|
|
"type": "txt"
|
|
, "content": self.load_txt_data(file_path)
|
|
})
|
|
return data
|
|
@staticmethod
|
|
def load_json_data(file_path):
|
|
with open(file_path) as file:
|
|
return json.load(file)
|
|
@staticmethod
|
|
def load_txt_data(file_path):
|
|
with open(file_path) as file:
|
|
return file.read() |