| import os.path |
| import json |
| from langchain.docstore.document import Document |
| from langchain.text_splitter import RecursiveCharacterTextSplitter |
| import ollama |
| from tqdm import tqdm |
|
|
| from doc_processor import DocProcessor |
|
|
| class ContextualDocProcessor(DocProcessor): |
| def __init__(self,DATA_DIR, PATH_SAVE, PATH_SAVE_CONTEXT): |
| DocProcessor.__init__(self,DATA_DIR, PATH_SAVE) |
| self.PATH_SAVE_CONTEXT = PATH_SAVE_CONTEXT |
|
|
| def process_data(self): |
| DocProcessor.process_data(self) |
| if os.path.exists(self.PATH_SAVE_CONTEXT): |
| self.load_context_from_jsonl() |
| else: |
| self.create_context() |
| self.save_context_to_jsonl() |
|
|
|
|
| def create_context(self): |
| self.context = [] |
| for filename in tqdm(os.listdir(self.DATA_DIR)): |
| if filename.endswith(".json"): |
| path = os.path.join(self.DATA_DIR, filename) |
| with open(path, 'r') as f: |
| doc = str(json.load(f)) |
| prompt = f"Résume le document ci-dessous en 500 caractères maximum:\n\nDocument:\n{doc}" |
| response = ollama.chat( |
| model="llama3", |
| messages=[ |
| { |
| "role": "user", |
| "content": prompt, |
| }, |
| ], |
| ) |
|
|
| self.context += [Document(page_content=response["message"]["content"], metadata={"source": path}, id=path)] |
|
|
| print("Chunks created") |
|
|
| def save_context_to_jsonl(self): |
| with open(self.PATH_SAVE_CONTEXT, 'w') as jsonl_file: |
| for doc in self.context: |
| jsonl_file.write(doc.json() + '\n') |
| print("Context saved") |
|
|
| def load_context_from_jsonl(self): |
| self.context = [] |
| with open(self.PATH_SAVE_CONTEXT, 'r') as jsonl_file: |
| for line in jsonl_file: |
| data = json.loads(line) |
| obj = Document(**data) |
| self.context.append(obj) |
| print("Context loaded") |
| |
|
|
|
|
| |
