Loads data from various file types and chunks it into an ensemble retriever.
(self)
| 144 | ### Data Loading Section |
| 145 | ######################## |
| 146 | def load_data(self): |
| 147 | """ |
| 148 | Loads data from various file types and chunks it into an ensemble retriever. |
| 149 | """ |
| 150 | data_dir = os.getenv("data_directory") |
| 151 | file_types = os.getenv("file_types").split(",") |
| 152 | |
| 153 | if "json" in file_types: |
| 154 | jq_compiled = jq.compile(os.getenv("json_schema")) |
| 155 | |
| 156 | # Load all files in the data directory, recursively |
| 157 | files = glob.glob(os.path.join(data_dir, "**"), recursive=True) |
| 158 | documents = [] |
| 159 | with tqdm( |
| 160 | total=len(files), desc="Reading in, chunking, and vectorizing documents" |
| 161 | ) as pbar: |
| 162 | for file in files: |
| 163 | file_type = file.split(".")[-1] |
| 164 | if os.path.isfile(file) and file_type in file_types: |
| 165 | # Load the document based on the file type |
| 166 | if file_type == "json": |
| 167 | with open(file, "r", encoding="utf-8") as f: |
| 168 | doc = json.load(f) |
| 169 | doc = jq_compiled.input(doc).first() |
| 170 | doc = json.dumps(doc) |
| 171 | elif file_type == "txt" or file_type == "xml": |
| 172 | with open(file, "r", encoding="utf-8") as f: |
| 173 | doc = f.read() |
| 174 | elif file_type == "csv": |
| 175 | import pandas as pd |
| 176 | df = pd.read_csv(file, encoding="utf-8", sep=os.getenv("csv_separator")) |
| 177 | json_data = df.to_dict(orient='records') |
| 178 | doc = json.dumps(json_data) |
| 179 | elif file_type == "pptx": |
| 180 | presentation = Presentation(file) |
| 181 | full_text = [] |
| 182 | for slide in presentation.slides: |
| 183 | slide_text = [] |
| 184 | for shape in slide.shapes: |
| 185 | if shape.has_text_frame: |
| 186 | for paragraph in shape.text_frame.paragraphs: |
| 187 | slide_text.append(paragraph.text) |
| 188 | full_text.append("\n".join(slide_text)) |
| 189 | doc = "\n\n".join(full_text) |
| 190 | else: |
| 191 | doc = self.converter.convert(file).document.export_to_text() |
| 192 | |
| 193 | # Get the subfolder name of this document |
| 194 | subfolder = os.path.basename(os.path.dirname(file)).replace(os.getenv("data_directory"), "") |
| 195 | if not(file_type == "csv"): |
| 196 | # Chunk the document |
| 197 | chunks = self.splitter.split_text(doc) |
| 198 | else: |
| 199 | chunks = doc |
| 200 | |
| 201 | chunks = [{ |
| 202 | "id": hashlib.md5(chunk.encode()).hexdigest(), |
| 203 | "embedding": self.embeddings.encode(chunk), |
no test coverage detected