Fine-Tuning a Legal Copilot Using Azure OpenAI and W&B

from datasets import load_dataset
import json

# List of subsets to load
subsets = [
    "contract_nli_confidentiality_of_agreement",
    "contract_nli_explicit_identification",
    "contract_nli_inclusion_of_verbally_conveyed_information",
    "contract_nli_limited_use",
    "contract_nli_no_licensing",
    "contract_nli_notice_on_compelled_disclosure",
    "contract_nli_permissible_acquirement_of_similar_information",
    "contract_nli_permissible_copy",
    "contract_nli_permissible_development_of_similar_information",
    "contract_nli_permissible_post-agreement_possession",
    "contract_nli_return_of_confidential_information",
    "contract_nli_sharing_with_employees",
    "contract_nli_sharing_with_third-parties",
    "contract_nli_survival_of_obligations",
    "contract_qa"
]

# Dictionary to hold subsets as keys and their prompts as values
subsets_prompts = {
    "contract_nli_confidentiality_of_agreement": "Identify if the clause provides that the Receiving Party shall not disclose the fact that Agreement was agreed or negotiated.",
    "contract_nli_explicit_identification": "Identify if the clause provides that all Confidential Information shall be expressly identified by the Disclosing Party.",
    "contract_nli_inclusion_of_verbally_conveyed_information": "Identify if the clause provides that Confidential Information may include verbally conveyed information.",
    "contract_nli_limited_use": "Identify if the clause provides that the Receiving Party shall not use any Confidential Information for any purpose other than the purposes stated in Agreement.",
    "contract_nli_no_licensing": "Identify if the clause provides that the Agreement shall not grant Receiving Party any right to Confidential Information.",
    "contract_nli_notice_on_compelled_disclosure": "Identify if the clause provides that the Receiving Party shall notify Disclosing Party in case Receiving Party is required by law, regulation or judicial process to disclose any Confidential Information.",
    "contract_nli_permissible_acquirement_of_similar_information": "Identify if the clause provides that the Receiving Party may acquire information similar to Confidential Information from a third party.",
    "contract_nli_permissible_copy": "Identify if the clause provides that the Receiving Party may create a copy of some Confidential Information in some circumstances.",
    "contract_nli_permissible_development_of_similar_information": "Identify if the clause provides that the Receiving Party may independently develop information similar to Confidential Information.",
    "contract_nli_permissible_post-agreement_possession": "Identify if the clause provides that the Receiving Party may retain some Confidential Information even after the return or destruction of Confidential Information.",
    "contract_nli_return_of_confidential_information": "Identify if the clause provides that the Receiving Party shall destroy or return some Confidential Information upon the termination of Agreement.",
    "contract_nli_sharing_with_employees": "Identify if the clause provides that the Receiving Party may share some Confidential Information with some of Receiving Party's employees.",
    "contract_nli_sharing_with_third-parties": "Identify if the clause provides that the Receiving Party may share some Confidential Information with some third-parties (including consultants, agents and professional advisors).",
    "contract_nli_survival_of_obligations": "Identify if the clause provides that some obligations of Agreement may survive termination of Agreement.",
    "contract_qa": "Answer questions about whether contractual clauses discuss particular issues."
}


# Path to save the JSONL files
output_file_train = './legalbench_subsets_test.jsonl'
output_file_test = './legalbench_subsets_train.jsonl'

def write_to_file(subset, output_file, split):
    with open(output_file, 'a') as file:
        dataset = load_dataset("nguha/legalbench", subset, split=split)

        for example in dataset:
            # Constructing the dialogue format with the prompt and example data
            prompt = subsets_prompts.get(subset, "")
            dialogue = {
                "messages": [
                    {"role": "system", "content": "Answer yes/no "+ prompt},
                    {"role": "user", "content": example['text']},
                    {"role": "assistant", "content": example['answer']}
                ]
            }
            file.write(json.dumps(dialogue) + '\n')

# Write train and test datasets to separate files
for subset in subsets:
    # Write the test data to the train file (since they are mixed up on Hugging Face)
    write_to_file(subset, output_file_train, "test")

    # Write the train data to the test file
    write_to_file(subset, output_file_test, "train")

{
"messages": 
[{"role": "system", "content": "Answer yes/no Identify if the clause provides that Confidential Information may include verbally conveyed information."}, 
{"role": "user", "content": "2. Every contract party can disclose confidential information to the other contract party orally or in writing.  "},
 {"role": "assistant", "content": "Yes"}
]
}

import os 
from openai import AzureOpenAI

client = AzureOpenAI(
  azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT"),
  api_key=os.getenv("AZURE_OPENAI_KEY"),
  api_version="2023-12-01-preview"  # This API version or later is required to access fine-tuning for turbo/babbage-002/davinci-002
)

# File names for training and validation datasets
training_file_name = './legalbench_subsets_train.jsonl'
validation_file_name = './legalbench_subsets_val.jsonl'

training_response = client.files.create(
    file=open(training_file_name, "rb"), purpose="fine-tune"
)
training_file_id = training_response.id

validation_response = client.files.create(
    file=open(validation_file_name, "rb"), purpose="fine-tune"
)
validation_file_id = validation_response.id

print("Training file ID:", training_file_id)
print("Validation file ID:", validation_file_id)

response = client.fine_tuning.jobs.create(
    training_file=training_file_id,
    validation_file=validation_file_id,
    model="gpt-35-turbo-0613", # Enter base model name. Note that in Azure OpenAI the model name contains dashes and cannot contain dot/period characters.
)
job_id = response.id
# You can use the job ID to monitor the status of the fine-tuning job.
# The fine-tuning job will take some time to start and complete.
print("Job ID:", job_id)
print(response.model_dump_json(indent=2))

import os
import openai
from wandb.integration.openai.fine_tuning import WandbLogger
from openai import AzureOpenAI
import argparse

def main():
    # Parse command line arguments
    parser = argparse.ArgumentParser(description='Sync Wandb with OpenAI Fine-Tune')
    parser.add_argument('--id', type=str, help='Fine Tune Job ID')
    parser.add_argument('--project', type=str, help='Project Name')
    parser.add_argument('--overwrite', action='store_true', help='Overwrite existing logs (defaults to False)')
    args = parser.parse_args()

    # Retrieve OpenAI API key from environment variable
    openai_api_key = os.getenv("OPENAI_API_KEY")
    if not openai_api_key:
        raise ValueError("OPENAI_API_KEY environment variable not set")

    # Initialize OpenAI with the API key
    openai.api_key = openai_api_key

    # Initialize AzureOpenAI client (modify the endpoint if necessary)
    client = AzureOpenAI(
        azure_endpoint = "YOUR_AZURE_ENDPOINT_URL",  # Replace with your Azure endpoint URL
        api_key=openai_api_key,
        api_version="2023-12-01-preview" 
    )

    # Sync Wandb with the fine-tune job
    WandbLogger.sync(fine_tune_job_id=args.id, openai_client=client, overwrite=args.overwrite)

if __name__ == "__main__":
    main()

from flask import Flask, request, jsonify, send_from_directory
import PyPDF2
import io
import os


app = Flask(__name__)

# Set debug mode (set to True for testing without OpenAI calls)
DEBUG_MODE = False 

from openai import AzureOpenAI

client = AzureOpenAI(
  azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT"),
  api_key=os.getenv("AZURE_OPENAI_KEY"),
  api_version="2023-12-01-preview"  # This API version or later is required to access fine-tuning for turbo/babbage-002/davinci-002
)


def getCompletion(prompt):
    chat_completion = client.chat.completions.create(
        model="deploy_V1", 
        messages=[
                    {"role": "system", "content": "You are a helpful assistant reviewing legal documents."},
                        {"role": "user", "content": prompt}
                    ]
    )
    return chat_completion



@app.route('/')
def index():
    # Serve the HTML file
    return send_from_directory('', 'index.html')

@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({"error": "No file part"}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    if file:
        pdfReader = PyPDF2.PdfReader(io.BytesIO(file.read()))
        responses = []
        combined_response = ""

        for page_number, page in enumerate(pdfReader.pages):
            text = page.extract_text()
            if DEBUG_MODE:
                # In debug mode, return the first 10 characters of each page
                response_text = text[:10]
            else:
                # Constructing the prompt for the model
                prompt = f"Review this legal document excerpt for any elements that seem legally unfair under US law \"{text}\""

                response = getCompletion(prompt=prompt)             
                # Extract response from OpenAI
                response_text = response.choices[0].message.content

            # Collecting responses for each page
            combined_response += f"Page {page_number + 1}: {response_text}\n\n"

        return jsonify({"text": combined_response})

    return jsonify({"error": "Invalid file"}), 400

if __name__ == '__main__':
    app.run(debug=True)

        for page_number, page in enumerate(pdfReader.pages):
            text = page.extract_text()
            if DEBUG_MODE:
                # In debug mode, return the first 10 characters of each page
                response_text = text[:10]
            else:
                # Constructing the prompt for the model
                prompt = f"Review this legal document excerpt for any elements that seem legally unfair under US law \"{text}\""

                response = getCompletion(prompt=prompt)             
                # Extract response from OpenAI
                response_text = response.choices[0].message.content