File size: 12,186 Bytes
dbe9c08 335a37e dbe9c08 358cb7f dbe9c08 88c55d1 dbe9c08 a4d6114 dbe9c08 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 |
import os
import json
import nltk
import openai
import chromadb
from langchain.document_loaders import UnstructuredXMLLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.document_loaders import TextLoader
from langchain.prompts import PromptTemplate
from langchain.chains import AnalyzeDocumentChain
from langchain.chains.question_answering import load_qa_chain
from langchain.callbacks import get_openai_callback
from langchain.llms import OpenAI
from langchain.vectorstores import FAISS
from langchain.text_splitter import CharacterTextSplitter
# Clear ChromaDB cache to fix tenant issue
chromadb.api.client.SharedSystemClient.clear_system_cache()
# Move variables and functions that don't need to be in the main function outside
nltk.download("punkt", quiet=True)
from nltk import word_tokenize, sent_tokenize
openai.api_key = os.getenv("OPENAI_API_KEY")
if openai.api_key is None:
raise Exception("OPENAI_API_KEY not found in environment variables")
embeddings = OpenAIEmbeddings()
def split_docs(documents, chunk_size=1000, chunk_overlap=0):
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=chunk_size, chunk_overlap=chunk_overlap
)
return text_splitter.split_documents(documents)
def call_QA_to_json(
prompt, year, month, day, saved_patent_names, index=0, logging=True, model_name="gpt-3.5-turbo"
):
"""
Generate embeddings from txt documents, retrieve data based on the provided prompt, and return the result as a JSON object.
Parameters:
prompt (str): The input prompt for the retrieval process.
year (int): The year part of the data folder name.
month (int): The month part of the data folder name.
day (int): The day part of the data folder name.
saved_patent_names (list): A list of strings containing the names of saved patent text files.
index (int): The index of the saved patent text file to process. Default is 0.
logging (bool): The boolean to print logs
Returns:
tuple: A tuple containing two elements:
- Cost of OpenAI API
- A JSON string representing the output from the retrieval chain.
This function loads the specified txt file, generates embeddings from its content,
and uses a retrieval chain to retrieve data based on the provided prompt.
The retrieved data is returned as a JSON object, and the raw documents are returned as a list of strings.
The output is also written to a file in the 'output' directory with the name '{index}.json'.
"""
llm = ChatOpenAI(model_name=model_name, temperature=0, cache=False)
file_path = os.path.join(
os.getcwd(),
"data",
"ipa" + str(year)[2:] + f"{month:02d}" + f"{day:02d}",
saved_patent_names[index],
)
if logging:
print(f"Loading documents from: {file_path}")
loader = TextLoader(file_path)
documents_raw = loader.load()
documents = split_docs(documents_raw)
if logging:
print("Generating embeddings and persisting...")
vectordb = Chroma.from_documents(
documents=documents, embedding=embeddings,
)
# vectordb.persist()
PROMPT_FORMAT = """
Task: Use the following pieces of context to answer the question at the end.
{context}
Question: {question}
"""
PROMPT = PromptTemplate(
template=PROMPT_FORMAT, input_variables=["context", "question"]
)
chain_type_kwargs = {"prompt": PROMPT}
retrieval_chain = RetrievalQA.from_chain_type(
llm, chain_type="stuff",
retriever=vectordb.as_retriever(),
chain_type_kwargs=chain_type_kwargs,
# return_source_documents=True
)
if logging:
print("Running retrieval chain...")
with get_openai_callback() as cb:
output = retrieval_chain.run(prompt)
if logging:
print(f"Total Tokens: {cb.total_tokens}")
print(f"Prompt Tokens: {cb.prompt_tokens}")
print(f"Completion Tokens: {cb.completion_tokens}")
print(f"Successful Requests: {cb.successful_requests}")
print(f"Total Cost (USD): ${cb.total_cost}")
cost = cb.total_cost
try:
# Convert output to dictionary
output_dict = json.loads(output)
# Manually assign the Patent Identifier
output_dict["Patent Identifier"] = saved_patent_names[index].split("-")[0]
# Check if the directory 'output' exists, if not create it
if not os.path.exists("output"):
os.makedirs("output")
if logging:
print("Writing the output to a file...")
with open(f"output/{saved_patent_names[index]}_{model_name}.json", "w", encoding="utf-8") as json_file:
json.dump(output_dict, json_file, indent=4, ensure_ascii=False)
if logging:
print("Call to 'call_QA_to_json' completed.")
except Exception as e:
print("An error occurred while processing the output.")
print("Error message:", str(e))
try:
vectordb.delete(ids=["*"])
except Exception as e:
print(f"Error deleting vector database: {str(e)}")
return cost, output
def call_TA_to_json(
prompt, year, month, day, saved_patent_names, index=0, logging=True
):
"""
Retrieve text analytics (TA) data from a specified patent file and convert the output to JSON format.
This function reads a text document from the patent file specified by the year, month, day, and file name parameters.
It then applies a QA retrieval process to the document using the provided prompt.
The result of the QA retrieval process is converted to a JSON object, which is then written to a file.
Additionally, a patent identifier is manually assigned to the output JSON object.
Parameters:
prompt (str): The input prompt for the retrieval process.
year (int): The year part of the data folder name.
month (int): The month part of the data folder name.
day (int): The day part of the data folder name.
saved_patent_names (list): A list of strings containing the names of saved patent text files.
index (int, optional): The index of the saved patent text file to process. Default is 0.
logging (bool, optional): If True, print logs to the console. Default is True.
Returns:
tuple: A tuple containing two elements:
- documents_raw (str): The raw document content loaded from the specified patent file.
- output (str): A JSON string representing the output from the TA retrieval process.
Note:
The output is also written to a file in the 'output' directory with the same name as the input file and a '.json' extension.
"""
llm = ChatOpenAI(model_name='gpt-3.5-turbo', cache=False)
file_path = os.path.join(
os.getcwd(),
"data",
"ipa" + str(year)[2:] + f"{month:02d}" + f"{day:02d}",
saved_patent_names[index],
)
if logging:
print(f"Loading documents from: {file_path}")
with open(file_path, 'r') as f:
documents_raw = f.read()
PROMPT_FORMAT = """
Task: Use the following pieces of context to answer the question at the end.
Question:
"""
prompt = PROMPT_FORMAT + prompt
qa_chain = load_qa_chain(llm, chain_type="map_reduce")
qa_document_chain = AnalyzeDocumentChain(combine_docs_chain=qa_chain)
if logging:
print("Running Analyze Document chain...")
output = qa_document_chain.run(input_document=documents_raw, question=prompt)
try:
# Convert output to dictionary
output_dict = json.loads(output)
# Manually assign the Patent Identifier
output_dict["Patent Identifier"] = saved_patent_names[index].split("-")[0]
# Check if the directory 'output' exists, if not create it
if not os.path.exists("output"):
os.makedirs("output")
if logging:
print("Writing the output to a file...")
# Write the output to a file in the 'output' directory
with open(f"output/{saved_patent_names[index]}.json", "w", encoding="utf-8") as json_file:
json.dump(output_dict, json_file, indent=4, ensure_ascii=False)
if logging:
print("Call to 'call_QA_to_json' completed.")
except Exception as e:
print("An error occurred while processing the output.")
print("Error message:", str(e))
return documents_raw, output
def call_QA_faiss_to_json(
prompt, year, month, day, saved_patent_names, index=0, logging=True, model_name="gpt-3.5-turbo"
):
"""
Generate embeddings from txt documents, retrieve data based on the provided prompt, and return the result as a JSON object.
Parameters:
prompt (str): The input prompt for the retrieval process.
year (int): The year part of the data folder name.
month (int): The month part of the data folder name.
day (int): The day part of the data folder name.
saved_patent_names (list): A list of strings containing the names of saved patent text files.
index (int): The index of the saved patent text file to process. Default is 0.
logging (bool): The boolean to print logs
Returns:
tuple: A tuple containing two elements:
- A list of strings representing the raw documents loaded from the specified XML file.
- A JSON string representing the output from the retrieval chain.
This function loads the specified txt file, generates embeddings from its content,
and uses a retrieval chain to retrieve data based on the provided prompt.
The retrieved data is returned as a JSON object, and the raw documents are returned as a list of strings.
The output is also written to a file in the 'output' directory with the name '{count}.json'.
"""
llm = ChatOpenAI(model_name=model_name, cache=False)
chain = load_qa_chain(llm, chain_type="stuff")
file_path = os.path.join(
os.getcwd(),
"data",
"ipa" + str(year)[2:] + f"{month:02d}" + f"{day:02d}",
saved_patent_names[index],
)
if logging:
print(f"Loading documents from: {file_path}")
loader = TextLoader(file_path)
documents_raw = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=0)
documents = text_splitter.split_documents(documents_raw)
docsearch = FAISS.from_documents(documents, embeddings)
docs = docsearch.similarity_search(prompt)
if logging:
print("Running chain...")
with get_openai_callback() as cb:
output = chain.run(input_documents=docs, question=prompt)
print(f"Total Tokens: {cb.total_tokens}")
print(f"Prompt Tokens: {cb.prompt_tokens}")
print(f"Completion Tokens: {cb.completion_tokens}")
print(f"Successful Requests: {cb.successful_requests}")
print(f"Total Cost (USD): ${cb.total_cost}")
try:
# Convert output to dictionary
output_dict = json.loads(output)
# Manually assign the Patent Identifier
output_dict["Patent Identifier"] = saved_patent_names[index].split("-")[0]
# Check if the directory 'output' exists, if not create it
if not os.path.exists("output"):
os.makedirs("output")
if logging:
print("Writing the output to a file...")
# Write the output to a file in the 'output' directory
with open(f"output/{saved_patent_names[index]}_{model_name}.json", "w", encoding="utf-8") as json_file:
json.dump(output_dict, json_file, indent=4, ensure_ascii=False)
if logging:
print("Call to 'call_QA_to_json' completed.")
except Exception as e:
print("An error occurred while processing the output.")
print("Error message:", str(e))
docsearch.delete
return output |