doczyai-pipelines/fieldExtraction/src/tin_pull.py

import json
import boto3
import pandas as pd
import numpy as np
from datetime import datetime
import os
import anthropic
import re
from concurrent.futures import ThreadPoolExecutor, as_completed

import config
import utils

# Global configuration variables
S3_REGION = "us-east-2"
BEDROCK_REGION = "us-east-1"
LLM_SELECTED = "Claude Instant"
bedrock_runtime = config.BEDROCK_RUNTIME


def find_tin_numbers(text):
    # Regular expression to find sets of 9 digits with an optional dash after the first two digits
    pattern = r"\b\d{2}-?\d{7}\b"

    # Find all matches in the provided text
    matches = re.findall(pattern, text)

    return matches


def invoke_llm(context, llm_selected=LLM_SELECTED):
    prompt_data = f"""
    
    Human: Use the following pieces of context to provide a concise answer to the questions at the end. You must answer in python list format.
    
    {context}
    
    Question: What is the Group provider's taxpayer identification number stated in the contract? This is a 9 digit long number typically with a hyphen after the first 2 digits. Return only this 9 digit number with a hyphen after the first 2 digits. Do not elaborate or add context.
    
    Assistant: Answer in list format: ["""

    if llm_selected == "Claude 3 - Sonnet":
        model_id = "anthropic.claude-3-5-sonnet-20240620-v1:0"
        body = json.dumps(
            {
                "anthropic_version": "bedrock-2023-05-31",
                "max_tokens": 4096,
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": anthropic.HUMAN_PROMPT
                                + prompt_data
                                + anthropic.AI_PROMPT,
                            }
                        ],
                    }
                ],
                "temperature": 0.0,
            }
        )
    elif llm_selected in ["Claude Instant"]:
        model_id = "anthropic.claude-instant-v1"
        body = json.dumps(
            {
                "prompt": anthropic.HUMAN_PROMPT + prompt_data + anthropic.AI_PROMPT,
                "max_tokens_to_sample": 2048,
                "temperature": 0.0,
                "top_p": 1,
                "top_k": 250,
                "stop_sequences": [anthropic.HUMAN_PROMPT],
            }
        )
    else:
        raise ValueError("Unsupported LLM selected")

    try:
        response = bedrock_runtime.invoke_model(
            body=body,
            modelId=model_id,
            accept="application/json",
            contentType="application/json",
        )
        response_body = json.loads(response.get("body").read())
        if llm_selected == "Claude 3 - Sonnet":
            response_text = response_body["content"][0]["text"]
        elif llm_selected in ["Claude Instant"]:
            response_text = response_body["completion"]
        else:
            print("Unsupported LLM selected")
    except Exception as e:
        print(f"Error invoking LLM: {e}")
        response_text = "failed"
    return response_text.strip()


def main(filename, contract_text):
    print(f"Processing contract: {filename}")
    try:
        df = pd.DataFrame(columns=["Contract Name", "Raw value", "Final value"])

        tin_number = find_tin_numbers(contract_text)
        final_tin = list(set(tin_number))

        if len(final_tin) != 1:
            final_tin = find_tin_numbers(
                invoke_llm(contract_text, llm_selected=LLM_SELECTED)
            )

        df["Raw value"] = [tin_number]
        df["Final value"] = [final_tin]
        df["Contract Name"] = [filename]

        df.to_csv(f"tin_output/{filename}.csv", index=False)
    except Exception as e:
        print(f"Error processing contract: {e}")


# bucket = "centene-national-contracting-files"
# contract_path = "no_tins_text_files"

# s3_client = boto3.Session(aws_access_key_id=config.S3_ACCESS_KEY_ID,
#                           aws_secret_access_key=config.S3_SECRET_ACCESS_KEY,
#                           aws_session_token=config.S3_SESSION_TOKEN, region_name="us-east-2").client('s3')

# objects = s3_client.list_objects_v2(Bucket=bucket, Prefix=contract_path)
file_list = []

OUTPUT_DIRECTORY = r"tin_output/no_tins_text_files"
NUM_WORKERS = 10
BATCH_SIZE = 10


def process_file(item):
    filename, contract_text = item
    if filename.endswith(".txt") and not os.path.exists(
        os.path.join(OUTPUT_DIRECTORY, (filename[:-3] + "csv"))
    ):
        main(filename, contract_text)


def process_directory(input_dict):
    # for item in input_dict.items():
    with ThreadPoolExecutor(max_workers=NUM_WORKERS) as executor:
        futures = {executor.submit(process_file, item) for item in input_dict.items()}
        for future in as_completed(futures):
            future.result()  # This will raise an exception if the callable raised one


input_dict = utils.read_input(path="data/cnc/no_tins")
process_directory(input_dict)

individual_path = r"tin_output/no_tins_text_files"


def individual_to_dict(filename):
    df = pd.read_csv(os.path.join(individual_path, filename))
    df = df.dropna(axis=1, how="all")
    column_names = ["Filename", "Answer", "Final Answer"]
    new_column_names = column_names[: len(df.columns)] + list(
        df.columns[len(column_names) :]
    )
    df.columns = new_column_names
    df_dict = {"Filename": filename}
    df_dict["Answer"] = df["Answer"].iloc[0]
    df_dict["Final Answer"] = df["Final Answer"].iloc[0]
    return df_dict


all_dicts = []
for filename in os.listdir(individual_path):
    df_dict = individual_to_dict(filename)
    all_dicts.append(df_dict)
final_df = pd.DataFrame(all_dicts)
final_df.to_csv("results.csv", index=False)