Spaces:
Running
Running
| from constants import DOCUMENT_COLLECTION | |
| from openai_constants import ENTITY_EXTRACTION_PROMPT, ENTITY_EXTRACTION_FUNCTION, GPT35_PRICING | |
| def extract_all_documents(openai_instance, chunks): | |
| all_entities = {} | |
| all_usage = {} | |
| total_prompt_tokens = 0 | |
| total_completion_tokens = 0 | |
| print(f"Number of chunks to process :: {len(chunks)}") | |
| for chunk_idx, chunk in enumerate(chunks): | |
| print(f"Sending request to OpenAI for {chunk_idx}") | |
| openai_entities_out = openai_instance.generate_response(ENTITY_EXTRACTION_PROMPT, chunk, ENTITY_EXTRACTION_FUNCTION) | |
| print("OpenAI out received") | |
| print(openai_entities_out['function_output']) | |
| #for ent in openai_entities_out['function_output'].items(): | |
| for key, val in openai_entities_out['function_output'].items(): | |
| print(key, val) | |
| if key in all_entities: | |
| if isinstance(val, list): | |
| all_entities[key].extend(val) # Extend the existing list with the new list | |
| else: | |
| all_entities[key].append(val) # Append the value to the existing list | |
| else: | |
| if isinstance(val, list): | |
| all_entities[key] = val # Initialize the key with the list | |
| else: | |
| all_entities[key] = [val] | |
| if 'prompt_tokens' in openai_entities_out['usage']: | |
| total_prompt_tokens += openai_entities_out['usage']['prompt_tokens'] | |
| if 'completion_tokens' in openai_entities_out['usage']: | |
| total_completion_tokens += openai_entities_out['usage']['completion_tokens'] | |
| all_usage = { | |
| 'prompt_tokens':total_prompt_tokens, | |
| 'completion_tokens':total_completion_tokens, | |
| 'output_pricing': total_completion_tokens/1000 * GPT35_PRICING['input'], | |
| 'input_pricing':total_prompt_tokens/1000 * GPT35_PRICING['output'] | |
| } | |
| return all_entities, all_usage | |
| def process_insurance_document(pii_instance, mongo_instance, openai_instance, ocr_instance, | |
| document_path, document_id): | |
| print("---- \nInside Process insurance document function") | |
| ## save file to S3 | |
| document_s3_url = "" | |
| ## OCR | |
| try: | |
| #document_text = ocr_instance.extract_text_from_document(document_path) | |
| document_text = ocr_instance.llama_parse_ocr(document_path) | |
| ocr_status = "Completed" | |
| process_status = "OCR Completed" | |
| print(f"OCR complete") | |
| except Exception as ex: | |
| document_text = "" | |
| ocr_status = ex | |
| process_status = f"OCR Failed. {ex}" | |
| print(process_status) | |
| ## save ocr file to S3, add document S3 url | |
| ocr_document_s3_url = "" | |
| ## update ocr_status in db | |
| #mongo_instance.update(DOCUMENT_COLLECTION, | |
| # {'document_id':document_id}, | |
| # {'set':{'ocr_status':ocr_status, 'document_s3_url':document_s3_url, | |
| # 'ocr_document_s3_url':ocr_document_s3_url, 'process_status':process_status}}) | |
| print(f"OCR status updated in db") | |
| ## PII entity extraction and masking | |
| pii_entities = pii_instance.identify(document_text) | |
| print(f"pii entiites are :: {pii_entities}") | |
| pii_entities = pii_instance.add_mask(pii_entities) | |
| print(f"\npii_entities after adding mask :: {pii_entities}") | |
| masked_text = pii_instance.anonymize(pii_entities, document_text) | |
| print(f"\nPII anonumized text is :: {masked_text}") | |
| print(f"\nPII complete") | |
| ## Openai extraction | |
| chunks = ocr_instance.chunk_document(masked_text) | |
| openai_entities, all_usage = extract_all_documents(openai_instance, chunks) | |
| entity_extraction_status = 'Completed' | |
| process_status = 'Document term extraction completed' | |
| """try: | |
| openai_entities, all_usage = extract_all_documents(openai_instance, chunks) | |
| entity_extraction_status = 'Completed' | |
| process_status = 'Document term extraction completed' | |
| except Exception as ex: | |
| openai_entities = {} | |
| all_usage = {} | |
| entity_extraction_status = ex | |
| process_status = f"Document term extraction failed. {ex}" | |
| """ | |
| #openai_entities_out = { | |
| # 'status':"Success", | |
| # 'function_output':{}, | |
| # 'usage':{} | |
| #} | |
| print(f"openai_entities are :: {openai_entities}") | |
| print(f"Request to OpenAI complete") | |
| print("----------- \nProcessing complete\n ") | |
| ## Unmask PII entities in openai entities | |
| ## update entity extraction status in db | |
| #mongo_instance.update(DOCUMENT_COLLECTION, | |
| # {'document_id':document_id}, | |
| # {'set':{'entity_extraction_status':entity_extraction_status, | |
| # 'entities':openai_entities, 'process_status':process_status}}) | |
| #print(f"Entities updated in DB") | |
| out = { | |
| "entities":openai_entities, | |
| "masked_text":masked_text, | |
| "masked_entities":pii_entities | |
| } | |
| return out |