scripts/LLM_test_old.py

# python -m scripts.LLM_test  --gif test_GPT_boxes --episodes 1 --max-steps 8 --model text-davinci-003 --env-args size 6 --env-name SocialAI-ColorBoxesLLMCSParamEnv-v1 --in-context-path llm_data/in_context_color_boxes.txt
# python -m scripts.LLM_test  --gif test_GPT_asoc --episodes 1 --max-steps 8 --model text-ada-001 --env-args size 6  --env-name SocialAI-AsocialBoxInformationSeekingParamEnv-v1 --in-context-path llm_data/in_context_asocial_box.txt --feed-full-ep

# python -m scripts.LLM_test  --gif test_GPT_boxes --episodes 1 --max-steps 8 --model bloom_560m --env-args size 6 --env-name SocialAI-ColorBoxesLLMCSParamEnv-v1 --in-context-path llm_data/in_context_color_boxes.txt
# python -m scripts.LLM_test  --gif test_GPT_asoc --episodes 1 --max-steps 8 --model bloom_560m --env-args size 6  --env-name SocialAI-AsocialBoxInformationSeekingParamEnv-v1 --in-context-path llm_data/in_context_asocial_box.txt --feed-full-ep

## bloom 560m
# boxes
# python -m scripts.LLM_test --log llm_log/bloom_560m_boxes_no_hist  --gif evaluation --episodes 20 --max-steps 10 --model bloom_560m --env-args size 6 --env-name SocialAI-ColorBoxesLLMCSParamEnv-v1 --in-context-path llm_data/in_context_color_boxes.txt

# asocial
# python -m scripts.LLM_test --log llm_log/bloom_560m_asocial_no_hist   --gif evaluation --episodes 20 --max-steps 10 --model bloom_560m --env-args size 6  --env-name SocialAI-AsocialBoxInformationSeekingParamEnv-v1 --in-context-path llm_data/in_context_asocial_box.txt

# random
# python -m scripts.LLM_test --log llm_log/random_boxes  --gif evaluation --episodes 20 --max-steps 10 --model random --env-args size 6 --env-name SocialAI-ColorBoxesLLMCSParamEnv-v1 --in-context-path llm_data/in_context_color_boxes.txt

import argparse
import json
import requests
import time
import warnings
from n_tokens import estimate_price

import numpy as np
import torch
from pathlib import Path

from utils.babyai_utils.baby_agent import load_agent
from utils import *
from models import *
import subprocess
import os

from matplotlib import pyplot as plt

from gym_minigrid.wrappers import *
from gym_minigrid.window import Window
from datetime import datetime

from imageio import mimsave

def prompt_preprocessor(llm_prompt):
    # remove peer observations
    lines = llm_prompt.split("\n")
    new_lines = []
    for line in lines:
        if line.startswith("#"):
            continue

        elif line.startswith("Conversation"):
            continue

        elif "peer" in line:
            caretaker = True
            if caretaker:
                # show only the location of the caretaker

                # this is very ugly, todo: refactor this
                assert "there is a" in line
                start_index = line.index('there is a') + 11
                new_line = line[:start_index] + 'caretaker'

                new_lines.append(new_line)

            else:
                # no caretaker at all
                if line.startswith("Obs :") and "peer" in line:
                    # remove only the peer descriptions
                    line = "Obs :"
                    new_lines.append(line)
                else:
                    assert "peer" in line

        elif "Caretaker:" in line:
            # line = line.replace("Caretaker:", "Caretaker says: '") + "'"
            new_lines.append(line)

        else:
            new_lines.append(line)

    return "\n".join(new_lines)


# Parse arguments

parser = argparse.ArgumentParser()
parser.add_argument("--model", required=False,
                    help="text-ada-001")
parser.add_argument("--seed", type=int, default=0,
                    help="Seed of the first episode. The seed for the following episodes will be used in order: seed, seed + 1, ... seed + (n_episodes-1) (default: 0)")
parser.add_argument("--max-steps", type=int, default=5,
                    help="max num of steps")
parser.add_argument("--shift", type=int, default=0,
                    help="number of times the environment is reset at the beginning (default: 0)")
parser.add_argument("--argmax", action="store_true", default=False,
                    help="select the action with highest probability (default: False)")
parser.add_argument("--pause", type=float, default=0.5,
                    help="pause duration between two consequent actions of the agent (default: 0.5)")
parser.add_argument("--env-name", type=str,
                    # default="SocialAI-ELangColorBoxesTestInformationSeekingParamEnv-v1",
                    # default="SocialAI-AsocialBoxInformationSeekingParamEnv-v1",
                    default="SocialAI-ColorBoxesLLMCSParamEnv-v1",
                    required=False,
                    help="env name")
parser.add_argument("--in-context-path", type=str,
                    # default='llm_data/short_in_context_boxes.txt'
                    # default='llm_data/in_context_asocial_box.txt'
                    default='llm_data/in_context_color_boxes.txt',
                    required=False,
                    help="path to in context examples")
parser.add_argument("--gif", type=str, default="visualization",
                    help="store output as gif with the given filename", required=False)
parser.add_argument("--episodes", type=int, default=1,
                    help="number of episodes to visualize")
parser.add_argument("--env-args", nargs='*', default=None)
parser.add_argument("--agent_view", default=False, help="draw the agent sees (partially observable view)", action='store_true' )
parser.add_argument("--tile_size", type=int, help="size at which to render tiles", default=32 )
parser.add_argument("--mask-unobserved", default=False, help="mask cells that are not observed by the agent", action='store_true' )
parser.add_argument("--log", type=str, default="llm_log/episodes_log", help="log from the run", required=False)
parser.add_argument("--feed-full-ep", default=False, help="weather to append the whole episode to the prompt", action='store_true')
parser.add_argument("--skip-check", default=False, help="Don't estimate the price.", action="store_true")

args = parser.parse_args()

# Set seed for all randomness sources

seed(args.seed)

model = args.model


in_context_examples_path = args.in_context_path

print("env name:", args.env_name)
print("examples:", in_context_examples_path)
print("model:", args.model)

# datetime
now = datetime.now()
datetime_string = now.strftime("%d_%m_%Y_%H:%M:%S")
print(datetime_string)

# log filenames

log_folder = args.log+"_"+datetime_string+"/"
os.mkdir(log_folder)
evaluation_log_filename = log_folder+"evaluation_log.json"
prompt_log_filename = log_folder + "prompt_log.txt"
ep_h_log_filename = log_folder+"episode_history_query.txt"
gif_savename = log_folder + args.gif + ".gif"

assert "viz" not in gif_savename # don't use viz anymore


env_args = env_args_str_to_dict(args.env_args)
env = make_env(args.env_name, args.seed, env_args)

# env = gym.make(args.env_name, **env_args)
print(f"Environment {args.env_name} and args: {env_args_str_to_dict(args.env_args)}\n")

# Define agent
print("Agent loaded\n")

# prepare models

if args.model in ["text-davinci-003", "text-ada-001", "gpt-3.5-turbo-0301"]:
    import openai
    openai.api_key = os.getenv("OPENAI_API_KEY")

elif args.model in ["gpt2_large", "api_bloom"]:
    HF_TOKEN = os.getenv("HF_TOKEN")

elif args.model in ["bloom_560m"]:
    from transformers import BloomForCausalLM
    from transformers import BloomTokenizerFast

    hf_tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-560m", cache_dir=".cache/huggingface/")
    hf_model = BloomForCausalLM.from_pretrained("bigscience/bloom-560m", cache_dir=".cache/huggingface/")

elif args.model in ["bloom"]:
    from transformers import BloomForCausalLM
    from transformers import BloomTokenizerFast

    hf_tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom", cache_dir=".cache/huggingface/")
    hf_model = BloomForCausalLM.from_pretrained("bigscience/bloom", cache_dir=".cache/huggingface/")


def plt_2_rgb(env):
    # data = np.frombuffer(env.window.fig.canvas.tostring_rgb(), dtype=np.uint8)
    # data = data.reshape(env.window.fig.canvas.get_width_height()[::-1] + (3,))

    width, height = env.window.fig.get_size_inches() * env.window.fig.get_dpi()
    data = np.fromstring(env.window.fig.canvas.tostring_rgb(), dtype='uint8').reshape(int(height), int(width), 3)
    return data

def generate(text_input, model):
    # return "(a) move forward"
    if model == "dummy":
        print("dummy action forward")
        return "move forward"

    elif model == "random":
        print("random agent")
        return random.choice([
            "move forward",
            "turn left",
            "turn right",
            "toggle",
        ])

    elif model in ["gpt-3.5-turbo-0301"]:
        while True:
            try:
                c = openai.ChatCompletion.create(
                    model=model,
                    messages=[
                        # {"role": "system", "content": ""},
                        # {"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."},
                        # {"role": "user", "content": "Continue the following text in the most logical way.\n"+text_input}
                        {"role": "user", "content": text_input}
                    ],
                    max_tokens=3,
                    n=1,
                    temperature=0,
                    request_timeout=30,
                )
                break
            except Exception as e:
                print(e)
                print("Pausing")
                time.sleep(10)
                continue
        print("generation: ", c['choices'][0]['message']['content'])
        return c['choices'][0]['message']['content']

    elif model in ["text-davinci-003", "text-ada-001"]:
        while True:
            try:
                response = openai.Completion.create(
                    model=model,
                    prompt=text_input,
                    # temperature=0.7,
                    temperature=0.0,
                    max_tokens=3,
                    top_p=1,
                    frequency_penalty=0,
                    presence_penalty=0,
                    timeout=30
                )
                break

            except Exception as e:
                print(e)
                print("Pausing")
                time.sleep(10)
                continue

        choices = response["choices"]
        assert len(choices) == 1
        return choices[0]["text"].strip().lower()  # remove newline from the end

    elif model in ["gpt2_large", "api_bloom"]:
        # HF_TOKEN = os.getenv("HF_TOKEN")
        if model == "gpt2_large":
            API_URL = "https://api-inference.huggingface.co/models/gpt2-large"

        elif model == "api_bloom":
            API_URL = "https://api-inference.huggingface.co/models/bigscience/bloom"

        else:
            raise ValueError(f"Undefined model {model}.")

        headers = {"Authorization": f"Bearer {HF_TOKEN}"}

        def query(text_prompt, n_tokens=3):

            input = text_prompt

            # make n_tokens request and append the output each time - one request generates one token

            for _ in range(n_tokens):
                # prepare request
                payload = {
                    "inputs": input,
                    "parameters": {
                        "do_sample": False,
                        'temperature': 0,
                        'wait_for_model': True,
                        # "max_length": 500,  # for gpt2
                        # "max_new_tokens": 250  # fot gpt2-xl
                    },
                }
                data = json.dumps(payload)

                # request
                response = requests.request("POST", API_URL, headers=headers, data=data)
                response_json = json.loads(response.content.decode("utf-8"))

                if type(response_json) is list and len(response_json) == 1:
                    # generated_text contains the input + the response
                    response_full_text = response_json[0]['generated_text']

                    # we use this as the next input
                    input = response_full_text

                else:
                    print("Invalid request to huggingface api")
                    from IPython import embed; embed()

            # remove the prompt from the beginning
            assert response_full_text.startswith(text_prompt)
            response_text = response_full_text[len(text_prompt):]

            return response_text

        response = query(text_input).strip().lower()
        return response

    elif model in ["bloom_560m"]:
        # from transformers import BloomForCausalLM
        # from transformers import BloomTokenizerFast
        #
        # tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-560m", cache_dir=".cache/huggingface/")
        # model = BloomForCausalLM.from_pretrained("bigscience/bloom-560m", cache_dir=".cache/huggingface/")

        inputs = hf_tokenizer(text_input, return_tensors="pt")
        # 3 words
        result_length = inputs['input_ids'].shape[-1]+3
        full_output = hf_tokenizer.decode(hf_model.generate(inputs["input_ids"], max_length=result_length)[0])

        assert full_output.startswith(text_input)
        response = full_output[len(text_input):]

        response = response.strip().lower()

        return response

    else:
        raise ValueError("Unknown model.")

def get_parsed_action(text_action):
    if "move forward" in text_action:
        return "move forward"

    elif "turn left" in text_action:
        return "turn left"

    elif "turn right" in text_action:
        return "turn right"

    elif "toggle" in text_action:
        return "toggle"

    elif "no_op" in text_action:
        return "no_op"
    else:
        warnings.warn(f"Undefined action {text_action}")
        return "no_op"


def step(text_action):
    text_action = get_parsed_action(text_action)

    if "move forward" == text_action:
        action = [int(env.actions.forward), np.nan, np.nan]

    elif "turn left" == text_action:
        action = [int(env.actions.left), np.nan, np.nan]

    elif "turn right" == text_action:
        action = [int(env.actions.right), np.nan, np.nan]

    elif "toggle" == text_action:
        action = [int(env.actions.toggle), np.nan, np.nan]

    elif "no_op" == text_action:
        action = [np.nan, np.nan, np.nan]

    # if text_action.startswith("a"):
    #     action = [int(env.actions.forward), np.nan, np.nan]
    #
    # elif text_action.startswith("b"):
    #     action = [int(env.actions.left), np.nan, np.nan]
    #
    # elif text_action.startswith("c"):
    #     action = [int(env.actions.right), np.nan, np.nan]
    #
    # elif text_action.startswith("d"):
    #     action = [int(env.actions.toggle), np.nan, np.nan]
    #
    # elif text_action.startswith("e"):
    #     action = [np.nan, np.nan, np.nan]
    #
    # else:
    #     print("Unknown action.")

    obs, reward, done, info = env.step(action)

    return obs, reward, done, info



def reset(env):
    env.reset()
    # a dirty trick just to get obs and info
    return step("no_op")


def generate_text_obs(obs, info):
    llm_prompt = "Obs : "
    llm_prompt += "".join(info["descriptions"])
    if obs["utterance_history"] != "Conversation: \n":
        utt_hist = obs['utterance_history']
        utt_hist = utt_hist.replace("Conversation: \n","")
        llm_prompt += utt_hist

    return llm_prompt


def action_query():
    # llm_prompt = ""
    # llm_prompt += "Your possible actions are:\n"
    # llm_prompt += "(a) move forward\n"
    # llm_prompt += "(b) turn left\n"
    # llm_prompt += "(c) turn right\n"
    # llm_prompt += "(d) toggle\n"
    # llm_prompt += "(e) no_op\n"
    # llm_prompt += "Your next action is: ("
    llm_prompt = "Act :"
    return llm_prompt

# lod context examples
with open(in_context_examples_path, "r") as f:
    in_context_examples = f.read()

with open(prompt_log_filename, "a+") as f:
    f.write(datetime_string)

with open(ep_h_log_filename, "a+") as f:
    f.write(datetime_string)

feed_episode_history = args.feed_full_ep

# asoc
in_context_n_tokens = 800
ep_obs_len = 50 * 3

# color
in_context_n_tokens = 1434
# ep_obs_len = 70

# feed only current obs
if feed_episode_history:
    ep_obs_len = 50

else:
    # last_n = 1
    # last_n = 2
    last_n = 3
    ep_obs_len = 50 * last_n

_, price = estimate_price(
    num_of_episodes=args.episodes,
    in_context_len=in_context_n_tokens,
    ep_obs_len=ep_obs_len,
    n_steps=args.max_steps,
    model=args.model,
    feed_episode_history=feed_episode_history
)
if not args.skip_check:
    input(f"You will spend: {price} dollars. (in context: {in_context_n_tokens} obs: {ep_obs_len}), ok?")

# prepare frames list to save to gif
frames = []

assert args.max_steps <= 20

success_rates = []
# episodes start
for episode in range(args.episodes):
    print("Episode:", episode)
    new_episode_text = "New episode.\n"
    episode_history_text = new_episode_text

    success = False
    episode_seed = args.seed + episode
    env = make_env(args.env_name, episode_seed, env_args)

    with open(prompt_log_filename, "a+") as f:
        f.write("\n\n")

    observations = []
    actions = []
    for i in range(int(args.max_steps)):
        if i == 0:
            obs, reward, done, info = reset(env)
            action_text = ""

        else:
            with open(prompt_log_filename, "a+") as f:
                f.write("\nnew prompt: -----------------------------------\n")
                f.write(llm_prompt)

            text_action = generate(llm_prompt, args.model)
            obs, reward, done, info = step(text_action)
            action_text = f"Act : {get_parsed_action(text_action)}\n"
            actions.append(action_text)

            print(action_text)

        text_obs = generate_text_obs(obs, info)
        observations.append(text_obs)
        print(prompt_preprocessor(text_obs))

        # feed the full episode history
        episode_history_text += prompt_preprocessor(action_text + text_obs)  # append to history of this episode

        if feed_episode_history:
            # feed full episode history
            llm_prompt = in_context_examples + episode_history_text + action_query()

        else:
            n = min(last_n, len(observations))
            obs = observations[-n:]
            act = (actions + [action_query()])[-n:]

            episode_text = "".join([o+a for o,a in zip(obs, act)])

            llm_prompt = in_context_examples + new_episode_text + episode_text

        llm_prompt = prompt_preprocessor(llm_prompt)


        # save the image
        env.render(mode="human")
        rgb_img = plt_2_rgb(env)
        frames.append(rgb_img)

        if env.current_env.box.blocked and not env.current_env.box.is_open:
            # target box is blocked -> apple can't be obtained
            # break to save compute
            break

        if done:
            # quadruple last frame to pause between episodes
            for i in range(3):
                same_img = np.copy(rgb_img)
                # toggle a pixel between frames to avoid cropping when going from gif to mp4
                same_img[0, 0, 2] = 0 if (i % 2) == 0 else 255
                frames.append(same_img)

            if reward > 0:
                print("Success!")
                episode_history_text += "Success!\n"
                success = True
            else:
                episode_history_text += "Failure!\n"

            with open(ep_h_log_filename, "a+") as f:
                f.write("\nnew prompt: -----------------------------------\n")
                f.write(episode_history_text)

            break

        else:
            with open(ep_h_log_filename, "a+") as f:
                f.write("\nnew prompt: -----------------------------------\n")
                f.write(episode_history_text)

    print(f"{'Success' if success else 'Failure'}")
    success_rates.append(success)

mean_success_rate =  np.mean(success_rates)
print("Success rate:", mean_success_rate)
print(f"Saving gif to {gif_savename}.")
mimsave(gif_savename, frames, duration=args.pause)

print("Done.")

log_data_dict = vars(args)
log_data_dict["success_rates"] = success_rates
log_data_dict["mean_success_rate"] = mean_success_rate

print("Evaluation log: ", evaluation_log_filename)
with open(evaluation_log_filename, "w") as f:
    f.write(json.dumps(log_data_dict))