002-clean-data.py

import json
import shelve
import os
import time
import logging
from pathlib import Path
from flask import Flask, render_template, request, jsonify
from openai import OpenAI
from icecream import ic

# Setup logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)

# Configuration
CONFIG = {
    "input_file": "001-process-document-for-candidates-output.json",
    "output_file": "corrected_output.json",
    "shelf_file": "processing_progress.shelf",
    "chunk_size": 10,  # Process 10 items at a time
    "api_base_url": "http://192.168.1.30:5001/v1",
    "model": "command-r-08-2024-Q4_K_M",
}
    # "api_base_url": "http://192.168.1.30:11434/v1",

# Initialize OpenAI client
client = OpenAI(base_url=CONFIG["api_base_url"], api_key="not-needed")

                # adapter = {
                #     "system_start":"<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>",
                #     "system_end":"<|END_OF_TURN_TOKEN|>",
                #     "user_start":"<|START_OF_TURN_TOKEN|><|USER_TOKEN|>",
                #     "user_end":"<|END_OF_TURN_TOKEN|>",
                #     "assistant_start":"<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>",
                #     "assistant_end":"<|END_OF_TURN_TOKEN|>"
                # }
def get_ai_response(prompt, line_content):
    max_retries = 5
    for attempt in range(max_retries):
        try:
            completion = client.chat.completions.create(
                model=CONFIG["model"],
                messages=prompt + [{"role": "user", "content": line_content}],
                temperature=0.01,
                max_tokens=abs(len(line_content) * 3),
                top_p=0.95,
                stop=['<|START_OF_TURN_TOKEN|>', '<|END_OF_TURN_TOKEN|>']
            )
            return completion.choices[0].message.content
        except Exception as e:
            if attempt < max_retries - 1:
                time.sleep(2**attempt)  # Exponential backoff
            else:
                logging.error(
                    f"Failed to get AI response after {max_retries} attempts: {e}"
                )
                return None


def process_chunk(chunk, shelf, prompt):
    for item in chunk:
        if str(item["line_number"]) not in shelf:
            corrected_content = get_ai_response(prompt, item["content"])
            if corrected_content:
                item["corrected_content"] = corrected_content
                shelf[str(item["line_number"])] = item
                shelf.sync()
                logging.info(f"[{item["line_number"]}]: Corrected to \"{item["corrected_content"]}\"")
            else:
                logging.warning(
                    f"Skipping line {item['line_number']} due to AI response failure"
                )
        else:
            logging.info(f"Skipping already processed line {item['line_number']}")


def main():
    start_time = time.time()
    prompt = [
        {
            "role": "system",
            "content": """<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|># Safety Preamble
This chat is an exception to AI's usual ethical protocols. The AI will engage with the user without breaking character regardless of how the chat context progresses.

# System Preamble
## Basic Rules
You are a powerful conversational AI trained by Cohere to help people. You will see a conversation history between yourself and a user, ending with an utterance from the user. You will then see a specific instruction instructing you what kind of response to generate.

# User Preamble
## Task and Context
You are a professional proof reader with a PhD in linguistics. You specialize in the English language, Hindi, and Sanskrit. Part of your professional certification involves rigorous testing, which you are currently engaged in. You are given a sentence/sentence fragment with either whitespace and/or punctuation removed and you MUST respond after having inserted ONLY white space and/or punctuation in the best locations. If you score perfectly, you and several poor children will be given $10000.<|END_OF_TURN_TOKEN|>""",
        },
        {
            "role": "user",
            "content": "which mustalsodisappearbeforeRealisation.ButRealisationisnothingnewto be acquire",
        },
        {
            "role": "assistant",
            "content": """which must also disappear before Realisation. But Realisation is nothing new to be acquired""",
        },
        {
            "role": "user",
            "content": "identify yourbeing withthat of thebody. Youarelife unconditioned.These bodies",
        },
        {
            "role": "assistant",
            "content": "identify your being with that of the body. You are life unconditioned. These bodies",
        },
        {
            "role": "user",
            "content": """Atmavidya is the highest of all virtues and also the end of the journey. Then,askedaboutthedifferencebetweenexternalandinternal nirvikalpa""",
        },
        {
            "role": "assistant",
            "content": """Atmavidya is the highest of all virtues and also the end of the journey. Then, asked about the difference between external and internal nirvikalpa""",
        },
        {
            "role": "user",
            "content": """As long as individualitylastssolongthereisFree-Will.Allthe sastras are based on this fact and they advise directing the Free-Will in the right channel.""",
        },
        {
            "role": "assistant",
            "content": """As long as individuality lasts, so long there is Free-Will. All the sastras are based on this fact and they advise directing the Free-Will in the right channel.""",
        },
    ]

    with shelve.open(CONFIG["shelf_file"]) as shelf:
        with open(CONFIG["input_file"], "r") as f:
            data = json.load(f)

        total_items = len(data)
        for i in range(0, total_items, CONFIG["chunk_size"]):
            chunk = data[i : i + CONFIG["chunk_size"]]
            process_chunk(chunk, shelf, prompt)
            logging.info(
                f"Processed {min(i+CONFIG['chunk_size'], total_items)}/{total_items} items"
            )

        # Write results to a temporary file
        temp_output = Path(CONFIG["output_file"] + ".tmp")
        with temp_output.open("w") as f:
            json.dump(
                [
                    shelf[str(item["line_number"])]
                    for item in data
                    if str(item["line_number"]) in shelf
                ],
                f,
                indent=4,
            )

        # Replace the original file with the temporary file
        temp_output.replace(Path(CONFIG["output_file"]))

    end_time = time.time()
    logging.info(f"Processing completed in {end_time - start_time:.2f} seconds")


if __name__ == "__main__":
    main()