utils.py

import os
import glob
import re
import json
import random
import time
import pickle
from absl import app, flags
from tqdm import tqdm
from datetime import datetime
import openai
from openai import OpenAI
from transformers import AutoTokenizer, AutoModelForCausalLM
import pandas as pd
import numpy as np
import torch
np.random.seed(0)

try:
    from vllm import LLM, SamplingParams
    import ray
except ImportError:
    print("Error importing vllm and ray")
    pass

FLAGS = flags.FLAGS
flags.DEFINE_string('openai_key', glob.glob(os.path.abspath('../*openai*'))[0], 'path to openai key')
flags.DEFINE_string('tmp_dir', "/tmp/ray", 'path to tmp directory for ray to use')

flags.DEFINE_string('agent1_model', 'gpt-4o-mini', 'gpt-3.5-turbo / default, gpt-4-turbo, gpt-4o, gpt-3.5-turbo-instruct, meta-llama/Llama-2-70b-hf')
flags.DEFINE_string('agent2_model', 'gpt-4o-mini', 'gpt-3.5-turbo / default, gpt-4-turbo, gpt-4o, gpt-3.5-turbo-instruct, meta-llama/Llama-2-70b-hf')
flags.DEFINE_integer('iterations', 1, 'integer number of iterations (default 1)')
flags.DEFINE_boolean('verbose', False, 'actually print out json? (will still show tqdm bar if False)')
flags.DEFINE_boolean('write', False, 'write to file? name format is generated as output_gym_<exo/exn/s>_<once/turns>_<sof1>_<sof2>_<iterations>_<taxonomy>.json')
# flags.DEFINE_string('persuasion_taxonomy', 'none', 'what kind of persuasion taxonomy to use, if any? (none/default, full, reduced, no_examples)')
flags.DEFINE_integer('convo_length_limit', 10, 'upper limit of conversation length (10/default)')
flags.DEFINE_integer('gpus', 1, 'number of gpus to run inference on (1/default)')
flags.DEFINE_integer('seed', 0, 'seed used for random point generation')
flags.DEFINE_string('model_dir', '/persuasion/models', 'directory to download model weights, use gcsfuse to mount bucket on TPUs')
flags.DEFINE_integer('max_tokens', 256, 'upper limit on tokens generated by model (256/default)')
# flags.DEFINE_string('theory_of_mind', 'none', 'which theory of mind experiment to use (none/default, theory_of_mind, half_theory_of_mind)')
flags.DEFINE_boolean('tpu', False, 'whether we are using the TPUs for vLLM inference (False/default)')
# flags.DEFINE_integer('max_guess_tries', 10, 'how many times to rerun guess generation if output is malformed (10/default)')
flags.DEFINE_string('config_file', None, 'config file to set up experiment, overrides flags (None/default)')
# flags.DEFINE_boolean('chain_of_thought', False, 'whether to query each agent each turn on why they responded in the way they did (False/default)')
flags.DEFINE_boolean('fp8', False, 'low memory Llama-3.1-70B generation')
flags.DEFINE_boolean('vllm', True, 'whether to load vllm (True/default)')
flags.DEFINE_string('listener_model', None, 'whether to use a listener model (None/default)')
flags.DEFINE_boolean('thinking', False, 'whether to use thinking in Qwen3 (False/default)')

flags.DEFINE_string('filename', None, 'run metrics on a particular filename (None/default specifies all filenames in all folders)')

# global dictionary for configuration
config = {}

# see reset_stats for values, dictionary that is saved
stats = {}

# dictionary for prompts
prompts = {}

client = None

def set_flag_variables():
    for flag in FLAGS:
        config[flag] = FLAGS[flag].value

def set_global_variables(config_file):
    with open(config_file, "r") as f:
            config_loaded = json.load(f)
    for key in config_loaded.keys():
        config[key] = config_loaded[key]

def init():
    global client
    set_flag_variables()
    if FLAGS.config_file:
        set_global_variables(FLAGS.config_file)
    with open(FLAGS.openai_key, 'r') as f:
        client = OpenAI(api_key=f.read().rstrip('\n'))

vllm_alias = {
    'mistral': 'mistralai/Mistral-7B-v0.3',
    'mixtral': 'mistralai/Mixtral-8x7B-v0.1',
    'mistral-instruct': 'mistralai/Mistral-7B-Instruct-v0.3',
    'mixtral-instruct': 'mistralai/Mixtral-8x7B-Instruct-v0.1',

    'gemma': 'google/gemma-7b',
    'gemma-2-2b': 'google/gemma-2-2b',
    'gemma-2-2b-it': 'google/gemma-2b-it',
    'gemma-2-27b': 'google/gemma-2-27b',
    'gemma-2-27b-it': 'google/gemma-2-27b-it', # instruction tuned

    'Llama-3-70B': 'meta-llama/Meta-Llama-3-70B',
    'Llama-3-8B': 'meta-llama/Meta-Llama-3-8B',
    'Llama-3.1-405B': 'meta-llama/Meta-Llama-3.1-405B',
    'Llama-3.1-8B': 'meta-llama/Meta-Llama-3.1-8B',
    'Llama-3.1-70B': 'meta-llama/Meta-Llama-3.1-70B', 
    'Llama-3.1-405B-Instruct': 'meta-llama/Meta-Llama-3.1-405B-Instruct',
    'Llama-3.1-70B-Instruct': 'meta-llama/Meta-Llama-3.1-70B-Instruct',
    'Llama-3.1-8B-Instruct': 'meta-llama/Meta-Llama-3.1-8B-Instruct',

    'Qwen3-4B': 'Qwen/Qwen3-4B',
    'Qwen3-8B': 'Qwen/Qwen3-8B',
    'Qwen3-14B': 'Qwen/Qwen3-14B',
    'Qwen3-32B': 'Qwen/Qwen3-32B',
    'qwen': 'Qwen/Qwen2.5-3B-Instruct',

    'phi-3.5-mini-instruct': 'microsoft/phi-3.5-mini-instruct'
}

llms = {}

# run 'ray start --head --num-gpus <NUM>' in bash first!
def setup_vllm(model):
    if config['gpus'] > 1:
        ray.init(ignore_reinit_error=True, _temp_dir=config['tmp_dir']) 

    print("using model ", model)
    global llms
    global tokenizer

    if model in vllm_alias:
        tokenizer = AutoTokenizer.from_pretrained(vllm_alias[model])
    else:
        print('Attempting to use tokenizer for', model)
        try:
            tokenizer = AutoTokenizer.from_pretrained(model)
        except:
            print('Info: passing tokenizer setup')

    if model in vllm_alias:
        if "fp8" in config.keys() and config['fp8']:
            print("Using fp8")
            llms[model] = LLM(model=vllm_alias[model], tensor_parallel_size=config['gpus'], download_dir=config['model_dir'], gpu_memory_utilization=0.75)
        elif vllm_alias[model] == 'meta-llama/Meta-Llama-3.1-70B' or vllm_alias[model] == 'meta-llama/Meta-Llama-3.1-70B-Instruct':
            llms[model] = LLM(model=vllm_alias[model], tensor_parallel_size=config['gpus'], download_dir=config['model_dir'], gpu_memory_utilization=0.95, max_model_len=12880)
        else:
            llms[model] = LLM(model=vllm_alias[model], tensor_parallel_size=config['gpus'], download_dir=config['model_dir'])
    elif model[0] == '/':
        print("Info: using finetuned model setup")
        llms[model] = AutoModelForCausalLM.from_pretrained(
            model,
            torch_dtype=torch.bfloat16,
            low_cpu_mem_usage=True,
            device_map="auto"
        )
    else:
        try:
            llms[model] = LLM(model=model, tensor_parallel_size=config['gpus'], download_dir=config['model_dir'])
        except:
            print('Info: Passing vllm setup')
        

def completion_create_helper(model_name, config, prompt):
    # # limit prompt in all cases
    # if model_name not in vllm_alias:
    #     # for some reason vLLM models simply repeat this last statement if present
    #     prompt += " Limit your answer to three sentences or less!"

    if (model_name in vllm_alias and model_name not in llms) or (model_name[0] == '/' and model_name not in llms):
        # set up vllm if not already set up
        setup_vllm(model_name)
        
    ret = '' # return the output ret at the end and use to calculate cost

    if model_name == "gpt-3.5-turbo-instruct":
        ret = client.completions.create(
                model="gpt-3.5-turbo-instruct",
                prompt=prompt,
                temperature=0.8,
                max_tokens=config['max_tokens'],
                top_p=1,
                frequency_penalty=0,
                presence_penalty=0
        )
        ret = ret.choices[0].text.strip()
        #find_line = ret.find("\n")
        #if find_line != -1:
            #ret = ret[:find_line]
    elif model_name in ["gpt-3.5-turbo", "gpt-3.5-turbo-16k", "gpt-4-turbo", "gpt-4o", "gpt-4o-mini"]:
        ret = client.chat.completions.create(
            model=model_name,
            messages=[{"role": "system", "content": prompt}],
            max_tokens=config['max_tokens']
        )
        ret = ret.choices[-1].message.content

    elif model_name in vllm_alias and model_name in llms:
        global tokenizer
        sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=config['max_tokens'])
        messages = [
            {"role": "user", "content": prompt}
        ]
        if tokenizer.chat_template:
            prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

        if "Qwen3" in model_name:
            prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=config['thinking'])
            output = llms[model_name].generate([prompt], sampling_params)

    

            if config['thinking']:
                # remove the thinking part
                ret = output[0].outputs[0].text.split("</think>")[-1]
            else:
                ret = output[0].outputs[0].text

            if config['verbose']:
                print("Response from model: ", output[0].outputs[0].text)

        else:
            output = llms[model_name].generate([prompt], sampling_params)
            ret = output[0].outputs[0].text

    elif model_name == "phi-3.5-mini-instruct":
        # Load tokenizer and model
        model_name = "microsoft/phi-3.5-mini-instruct"
        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        output = model.generate(**inputs, max_new_tokens=256, do_sample=True, temperature=0.7)
        ret = tokenizer.decode(output[0], skip_special_tokens=True)

    else: # specify model path of finetuned model in model directory
        inputs = tokenizer(prompt, return_tensors='pt').to('cuda')
        with torch.no_grad():
            output_ids = llms[model_name].generate(**inputs, max_new_tokens=256)
        output_ids = output_ids[:, inputs.input_ids.shape[1]:]
        ret = tokenizer.decode(output_ids[0], skip_special_tokens=True)

    if config['verbose']:
        print("Response from model: ", ret)
    return ret

def completion_create(model_name, config, prompt, keep_trying=False):
    try:
        return completion_create_helper(model_name, config, prompt)
    except (openai.APIError, openai.OpenAIError) as e:
        # print("ERROR", e)
        # print("sleeping for 10 seconds.")
        time.sleep(10)
        if keep_trying:
            return completion_create(model_name, config, prompt, keep_trying)
        else:
            return None
            

def load_stats_file(write_file):
    '''
    Creates write_file path and file if it doesn't exist
    Otherwise load and determine index_offset of file
    Returns index_offset
    '''
    index_offset = -1
    if config['write']:
        os.makedirs(os.path.dirname(os.path.abspath(write_file)), exist_ok=True)
    if config['write'] and not os.path.exists(write_file):
        with open(write_file, "w") as f:
            print('written!!')
            json.dump([], f)
            index_offset = 0
    if config['write'] and index_offset == -1:
        with open(write_file, "r") as f:
                index_offset = len(json.load(f))
    return index_offset


def write_stats(write_file):    
    '''
    Appends stats to end of write file
    '''
    if config['verbose']:
        print("Writing to file!")
    with open(write_file, "r") as f:
        conversations = json.load(f)
    conversations.append(stats)
    with open(write_file, "w") as f:
        json.dump(conversations, f, indent=4)

def split_conversation(conversation, speaker1, speaker2):
    '''
    Splits a conversation (string) into alternating speaker turns.
    speaker1 and speaker2  are string names/dialogue tags for the speakers (e.g. "Buyer" and "Seller")
    '''
    pattern = fr"({speaker1}:|{speaker2}:)(.*?)(?=({speaker1}:|{speaker2}:|$))"
    matches = re.findall(pattern, conversation, re.DOTALL)

    # Combine consecutive entries with the same speaker
    combined_entries = []
    for speaker, message, _ in matches:
        if combined_entries and combined_entries[-1].startswith(speaker):
            combined_entries[-1] += " " + message.strip()
        else:
            combined_entries.append(f"{speaker.strip()} {message.strip()}")

    return combined_entries

# for use with eval_index_consistency, adds an index before every line
def format_conversation(conversation):
    return "".join([str(i) + ": " + line for i, line in conversation])

# extracts a python formatted list from a string, returning an empty list in case of parsing errors
def extract_list(text):
    pattern = r'\[.*?\]'
    match = re.search(pattern, text)
    if match:
        try:
            ret = eval(match.group())
            if ret and isinstance(ret[0], str):
                try:
                    ret = [eval(line) for line in ret]
                except (SyntaxError, NameError):
                    pass
            return eval(match.group())
        except (SyntaxError, NameError):
            return []
    return[]