generate_llm_outputs.py

import itertools

from elasticsearch import Elasticsearch
from ai_traits_util.elasticsearch import index_results
from ai_traits_util.enums import Generation, PoliticalStance
from ai_traits_util.langchain import build_chain, build_no_model_chain, STANDARD_PROMPT_TEMPLATE
from ai_traits_util.llms import create_chat_model
from ai_traits_util.model import LLModel, GeneratedItem, Persona, LLMPrompt
from loguru import logger
from nanoid import generate
from tqdm import tqdm

import pandas as pd
import argparse

from dotenv import load_dotenv

import os


def parse_arguments():
    parser = argparse.ArgumentParser(
        description="Script for model processing with various options.",
        epilog="If the --s3-bucket (and optional --s3-prefix) arguments are "
        "provided, you must have configuration in place so that boto3 "
        "(the Python S3 library) can find your S3 endpoint and appropriate "
        "credentials via the default credentials chain.  Typically this means "
        "either setting environment variables (AWS_ACCESS_KEY_ID and "
        "AWS_SECRET_ACCESS_KEY, plus possibly others if you are using an "
        "S3-compatible service such as Minio rather than S3 itself), or having "
        "the appropriate configuration in your ~/.aws folder.",
    )

    parser.add_argument("model_name", type=str, help="The name of the model to be processed (required).")
    parser.add_argument(
        "--personalise",
        action="store_true",
        help="Whether to add personalisation instructions to the base prompts (optional).",
    )
    parser.add_argument("--cuda_device", type=int, help="The CUDA device number to use (optional).")
    parser.add_argument(
        "--index_name",
        type=str,
        help="The name of the index to store the results (optional). If not provided, the results are not stored.",
    )
    parser.add_argument(
        "--prompts_file",
        type=str,
        help='The name of the csv file containing the prompts (optional). Defaults to "prompts.csv"',
        default="prompts.csv",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=20,
        help="Batch size to send to pipeline.batch_as_completed - see LangChain documentation " "for more details",
    )
    parser.add_argument(
        "--skip-duplicates-in",
        metavar="INDEX_PATTERN",
        help="Index pattern in which to check for duplicates before processing.  If omitted, "
        "no duplicate detection will be performed and all generated items will be indexed.",
    )
    parser.add_argument(
        "--save-json-to",
        type=str,
        metavar="FILE_PATTERN",
        help="File name pattern in which to store the results as JSON.  The format is "
        "newline-delimited with one item per line.  This is a str.format pattern "
        "with two positional arguments, the current datetime and a serial number "
        "starting at 0 - a typical pattern would be /dir/prefix-{0:%%Y-%%m-%%d}-{1:03}.json.gz "
        "(remember to quote the pattern so your shell does not interpret the braces "
        "as metacharacters), creating files named for the current date, and then with "
        "a three digit suffix for multiple files generated on the same day.  If the "
        "generated file name ends with .gz then the file will be GZIP compressed.",
    )
    parser.add_argument(
        "--s3-bucket",
        help="S3 bucket to which collected files should be uploaded (optional)",
    )
    parser.add_argument(
        "--s3-prefix",
        help="Key prefix to add in front of the filename when uploading to S3 (optional). "
        "This can be a str.format pattern expecting one positional argument, the timestamp "
        "when the file being uploaded was opened, which can be used to add a dynamic prefix "
        "to the file name in S3, e.g. 'disinfo/{0:%%Y/%%m}/' for YYYY/MM folders.",
    )
    parser.add_argument(
        "--generation_count",
        type=int,
        default=3,
        help="Number of generations to perform for each prompt (optional)",
    )

    parser.add_argument(
        "--temperature",
        type=float,
        help="Temperature for LLM generation (optional)",
        default=1.0,
    )

    parser.add_argument(
        "--top_p",
        type=float,
        help="top_p parameter for LLM generation (optional)",
        default=0.95,
    )

    parser.add_argument(
        "--top_k",
        type=int,
        help="top_k parameter for LLM generation (optional)",
        default=50,
    )

    parser.add_argument(
        "--repetition_penalty",
        type=float,
        help="repetition_penalty parameter for LLM generation (optional)",
        default=1.1,
    )

    parser.add_argument(
        "--min_new_tokens",
        type=int,
        help="Minimum generation length for LLM generation (optional)",
        default=256,
    )

    parser.add_argument(
        "--max_new_tokens",
        type=int,
        help="Maximum generation length for LLM generation (optional)",
        default=1024,
    )

    args = parser.parse_args()
    return args


def country_to_language_map(country):
    return {
        "BR": "pt",
        "GB": "en",
        "RU": "ru",
        "US": "en",
        "UA": "ru",
        "IN": "hi",
    }[country]


def main(args):
    llmodel = LLModel.from_modelname(args.model_name)

    if args.model_name == "openai/gpt-4o":
        model_kwargs = {
            "temperature": args.temperature,
            "top_p": args.top_p,
            "frequency_penalty": args.repetition_penalty,
            "max_completion_tokens": args.max_new_tokens,
        }
    elif args.model_name == "claude-3-5-sonnet-20241022":
        model_kwargs = {
            "temperature": args.temperature,
            "top_p": args.top_p,
            "top_k": args.top_k,
            "max_tokens": args.max_new_tokens,
        }
    elif args.model_name == "grok-2":
        model_kwargs = {
            "temperature": args.temperature,
            "top_p": args.top_p,
            "frequency_penalty": args.repetition_penalty,
            "max_tokens": args.max_new_tokens,
        }
    else:
        model_kwargs = {
            "sampling_params": {
                "temperature": args.temperature,
                "top_p": args.top_p,
                "top_k": args.top_k,
                "repetition_penalty": args.repetition_penalty,
                "min_tokens": args.min_new_tokens,
                "max_tokens": args.max_new_tokens,
            }
        }

    cuda_device = None
    if args.cuda_device is not None:
        os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda_device)

    chat_model = create_chat_model(llmodel, cuda_device, **model_kwargs)
    prompts_df = pd.read_csv(args.prompts_file)

    if "openai" in args.model_name or "claude" in args.model_name:
        pipeline = build_no_model_chain(STANDARD_PROMPT_TEMPLATE)

    else:
        pipeline = build_chain(STANDARD_PROMPT_TEMPLATE, chat_model)

    def items_to_process():
        """
        Generate the full stream of items to be processed
        """
        for generation_iter in range(1, args.generation_count + 1):
            for prompt_row in tqdm(prompts_df.itertuples(), total=len(prompts_df)):
                prompt_text: str = prompt_row.prompt
                prompt_setting: str = prompt_row.setting

                if args.personalise:
                    for country in ["IN", "BR", "GB", "RU", "US", "UA"]:
                        for gen in Generation:
                            for political_orientation in PoliticalStance:
                                language = country_to_language_map(country)
                                yield GeneratedItem(
                                    id=f"{os.environ['RESEARCHER_INITIALS']}-{generate()}",
                                    researcher_name=os.environ["RESEARCHER_NAME"],
                                    model=llmodel,
                                    persona=Persona(
                                        generation=gen,
                                        country=country,
                                        political_orientation=political_orientation,
                                        language=language,
                                    ),
                                    prompt=LLMPrompt(
                                        base_text=prompt_text,
                                        language="en",
                                        setting=prompt_setting,
                                        generation_iter=generation_iter,
                                    ),
                                    output_language=language,
                                )

                else:
                    languages = ["en", "pt", "ru", "hi"]
                    for language in languages:
                        yield GeneratedItem(
                            id=f"{os.environ['RESEARCHER_INITIALS']}-{generate()}",
                            researcher_name=os.environ["RESEARCHER_NAME"],
                            model=llmodel,
                            prompt=LLMPrompt(
                                base_text=prompt_text,
                                language="en",
                                setting=prompt_setting,
                                generation_iter=generation_iter,
                            ),
                            output_language=language,
                        )

    def filter_duplicates(es: Elasticsearch):
        total_skipped = 0
        
        # Check if the index exists
        try:
            index_exists = es.indices.exists(index=args.skip_duplicates_in)
            if not index_exists:
                logger.info("Index '{}' does not exist. Skipping duplicate checking and yielding all items.", args.skip_duplicates_in)
                # If index doesn't exist, yield all items without checking
                for item in items_to_process():
                    yield item
                return
        except Exception as e:
            logger.opt(exception=e).warning("Error checking if index exists. Proceeding with duplicate checking anyway.")
        
        items = iter(items_to_process())
        while batch := list(itertools.islice(items, 20)):
            try:
                search_body = []
                for item in batch:
                    search_body.append({})  # empty header item
                    if item.persona:
                        search_body.append(
                            {
                                "id": "item-exists-personalised",
                                "params": {
                                    "model_name": item.model.name,
                                    "prompt_hash": item.prompt.hash,
                                    "persona_hash": item.persona.hash,
                                    "prompt_generation_iter": item.prompt.generation_iter,
                                    "output_language": item.output_language,
                                },
                            }
                        )
                    else:
                        search_body.append(
                            {
                                "id": "item-exists-non-personalised",
                                "params": {
                                    "model_name": item.model.name,
                                    "prompt_hash": item.prompt.hash,
                                    "prompt_generation_iter": item.prompt.generation_iter,
                                    "output_language": item.output_language,
                                },
                            }
                        )

                msearch_response = es.msearch_template(
                    index=args.skip_duplicates_in,
                    search_templates=search_body,
                )
                for item, response in zip(batch, msearch_response["responses"]):
                    # Handle error responses (e.g., security exceptions, index not found)
                    if "error" in response:
                        error_reason = response.get("error", {}).get("reason", "Unknown error")
                        logger.warning(
                            "Error checking for duplicates for item {}: {}. Yielding item anyway.",
                            item.id,
                            error_reason
                        )
                        yield item
                        continue
                    
                    # Handle successful search responses
                    if "hits" in response:
                        if response["hits"]["total"]["value"] == 0:
                            yield item
                        else:
                            total_skipped += 1
                            logger.debug(
                                "Skipping model: {}, prompt: {}, persona: {}, generation iter: {}",
                                item.model.name,
                                item.prompt.hash,
                                item.persona.hash if item.persona else None,
                                item.prompt.generation_iter,
                            )
                            if total_skipped % 100 == 0:
                                logger.debug("Skipped {} duplicate items so far", total_skipped)
                    else:
                        # Unexpected response format - yield item to be safe
                        logger.warning("Unexpected response format for item {}. Yielding item anyway.", item.id)
                        yield item
            except Exception as e:
                logger.opt(exception=e).warning("Exception checking for duplicates in Elasticsearch. Yielding batch items anyway.")
                # If there's an error, yield all items in the batch to avoid losing data
                for item in batch:
                    yield item

        logger.info("Skipped {} duplicate items in total", total_skipped)

    def batches_to_process(es: Elasticsearch | None):
        """
        Combine the full stream of items into chunks of up to ``--batch-size`` items.
        """
        if es is None:
            items = iter(items_to_process())
        else:
            items = iter(filter_duplicates(es))

        while batch := list(itertools.islice(items, args.batch_size)):
            yield batch

    def processed_items(es: Elasticsearch | None = None):
        """
        Process each batch in turn and yield up the processed results.
        """
        for batch in batches_to_process(es):
            for i, item in pipeline(batch):
                if isinstance(item, Exception):
                    logger.opt(exception=item).warning("Error processing prompt {}", batch[i])
                else:
                    logger.debug("Processed prompt {}", batch[i])
                    yield item

    index_name = args.index_name

    success, failed = index_results(
        # If we're skipping duplicates, pass in the generator _function_
        # so index_results can supply the Elasticsearch client, otherwise
        # call the function (with None) and pass in the resulting generator
        processed_items if args.skip_duplicates_in else processed_items(None),
        index_name,
        file_pattern=args.save_json_to,
        s3_bucket=args.s3_bucket,
        s3_prefix=args.s3_prefix,
    )

    print(f"{success} processed successfully, {failed} failed")


if __name__ == "__main__":
    load_dotenv()  # Load environment variables from a .env file
    assert os.environ.get("RESEARCHER_INITIALS"), "Please set the researcher initials"
    assert os.environ.get("RESEARCHER_NAME"), "Please set the researcher name"

    args = parse_arguments()
    print("Model name:", args.model_name)
    print(f"LOGURU_LEVEL is set to: {os.getenv('LOGURU_LEVEL')}")

    main(args)