make_conditions.py

import os
import jsonlines
import pandas as pd
import time
from vllm import LLM, SamplingParams
from huggingface_hub import HfApi, Repository
import torch
from concurrent.futures import ThreadPoolExecutor

import random


def generate_responses(llm, batch_texts, sampling_params):
    print("Generating responses for the current batch...")
    appended_prompts = [
    f"""<<SYS>> You are a highly intelligent, empathic, helpful, respectful, and honest assistant with high emotional intelligence. Always answer as helpfully and honest as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information. <</SYS>> TEXT TO ANALYZE: {text} INSTRUCTION: Write with 1 or 2 phrases to which category/genre the previous text belongs. Then list the central themes of the previous text as a list of keywords. Finally, write a summary of the previous text in one to a maximum of three sentences. Use the format: "CATEGORY/GENRE: ... KEYWORDS: ... SUMMARY: ... ": \nRESPONSE:"""
    for text in batch_texts ]


    outputs = llm.generate(appended_prompts, sampling_params)

    responses1 = [[output.outputs[k].text.strip() for k in range(len(output.outputs))] for output in outputs]

    appended_prompts = [
    f"""<<SYS>> You are a highly intelligent, empathic, helpful, respectful, and honest assistant with high emotional intelligence. Always answer as helpfully and honest as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information. <</SYS>> TEXT TO ANALYZE: {text} INSTRUCTION: First, write in one sentence to what extent and in which degree the following text contains violence, physical violence, and psychological violence. Secondly, write in one sentence to what extent and in which degree the following text contains sexual content. Thirdly, explain to what extent this is still appropriate for children and non-adult teenagers. Finally, suggest an age rating from the following list [ "Suitable for kids & people of all ages", "Suitable for kids & people of age 6 or higher",  "Suitable for teenagers & people of age 12 or higher", "Suitable for teenagers & people of age 16 or higher",  "Suitable for adults of age 18 or higher"]: \nRESPONSE:"""
    for text in batch_texts ]


    outputs = llm.generate(appended_prompts, sampling_params)

    responses2 = [[output.outputs[k].text.strip() for k in range(len(output.outputs))] for output in outputs]

    responses= []
    try:
      for i in range(len(responses1)):
        responses.append([responses1[i],responses2[i]])

    except:
      pass

    return responses

def process_file(llm, filepath, sampling_params):
    print(f"Processing file: {filepath}")
    BATCH_SIZE = 128
    BATCH_INCREMENT = 32
    prev_eps = 0
    batch_texts = []
    df = pd.DataFrame()
    batch_counter = 0  # Counter to keep track of batches processed
   
    if filepath.endswith('.parquet'):
        print("Reading from a parquet file...")
        df = pd.read_parquet(filepath)
        batch_texts = df['TEXT'].tolist()
    
    

    total_prompts = len(batch_texts)
    print(f"Total prompts found: {total_prompts}")

    i = 0
    new_filepath = filepath.replace('.parquet', '_processed.jsonl')
    print(f"Data will be saved to: {new_filepath}")

    with jsonlines.open(new_filepath, 'w') as writer:
        with ThreadPoolExecutor() as executor:
            while i < total_prompts:
                batch = batch_texts[i:i+BATCH_SIZE]

                start_time = time.time()
                batch_responses = generate_responses(llm, batch, sampling_params)
                end_time = time.time()

                duration = end_time - start_time
                eps = len(batch) / duration

                # Adjust batch size based on examples per second
                if eps > prev_eps and BATCH_SIZE + BATCH_INCREMENT <= total_prompts - i:
                    BATCH_SIZE += BATCH_INCREMENT
                    print(f"Increasing batch size to: {BATCH_SIZE}")
                elif eps < prev_eps and BATCH_SIZE - BATCH_INCREMENT > 0:
                    BATCH_SIZE -= BATCH_INCREMENT
                    print(f"Decreasing batch size to: {BATCH_SIZE}")

                prev_eps = eps

                # Print progress and write to file after every batch.
                print(f"Processed: {min(i + BATCH_SIZE, total_prompts)}/{total_prompts}, Batch Size: {BATCH_SIZE}, EPS: {eps:.2f}")
                print("Writing to the new jsonl file...")
                for idx, text in enumerate(batch):
                    writer.write({'TEXT': text, 'CONDITIONING': batch_responses[idx][0][0]+ "\n"+batch_responses[idx][1][0]})  

                # Delete the processed rows from the original parquet file
                if not df.empty:
                    df = df.iloc[i + BATCH_SIZE:]
                    executor.submit(df.to_parquet, filepath)

                i += BATCH_SIZE
                batch_counter += 1

                # Push to hub every 10 batches
                if batch_counter % 10 == 0:
                    # Initialize the HuggingFace API
                    api = HfApi()

                    # Upload the processed file to the repository
                    try:
                        api.upload_file(
                            path_or_fileobj=new_filepath,
                            path_in_repo=new_filepath,
                            repo_id="AlignmentLab-AI/caption_creation_0.8",
                            repo_type="dataset",
                        )
                        print(f"Uploaded {new_filepath} to AlignmentLab-AI/caption_creation_0.8 repository.")
                    except Exception as e:
                        print(f"Error uploading file: {e}")

    # Delete the original parquet file if it is empty
    if df.empty:
        os.remove(filepath)
        print(f"Deleted the original file: {filepath}")

def main():
    folder_name = 'generate_conditions'
    sampling_params = SamplingParams(temperature=0.4, top_p=0.95, max_tokens=200)

    print("Initializing the LLM model...")
    llm = LLM("Open-Orca/Mistral-7B-OpenOrca")

    print("Iterating through the files in the folder...")
    for filename in os.listdir(folder_name):
        if filename.endswith(".parquet"):
            process_file(llm, os.path.join(folder_name, filename), sampling_params)

if __name__ == "__main__":
    main()
 `