Upload make_conditions.py (#2)

- Upload make_conditions.py (2ef50029118a0532c2332b8ec0f45f3f98b30647)


Co-authored-by: Miter Mix <mitermix@users.noreply.huggingface.co>

make_conditions.py

@@ -0,0 +1,138 @@
+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()
+ `