from utilities import *
from subprocess import STDOUT, check_output
from filelock import FileLock
import tempfile
import os
import func_timeout

"""2 dictionaries are maintained:
1. saved_evaluations: everything that we know is already saved to disk
2. buffered_evaluations: everything that we have computed but not yet saved to disk

Every 100 evaluations, we save the buffer to disk, then clear the buffer.
"""
saved_evaluations = None
buffered_evaluations = {}

def exec_and_return_output(code):
    OUTPUT = []
    exec(code)
    return OUTPUT

def cached_eval(expression, statement=False):
    global saved_evaluations, buffered_evaluations

    fn = "eval.pickle"
    if saved_evaluations is None:
        with FileLock(f'{fn}.lck'):
            if os.path.exists(fn):
                with open(fn, "rb") as handle:
                    saved_evaluations = pickle.load(handle)
            else:
                saved_evaluations = {}
        print("Loaded", len(saved_evaluations), "evaluations from disk")
    
    
    if expression in buffered_evaluations:
        result = buffered_evaluations[expression]
    elif expression in saved_evaluations: 
        result = saved_evaluations[expression]
    else:
        if statement:
            if False: # process sandbox
                # dump code to a temporary file stored on the RAM, /dev/shm
                # this is faster than writing to disk
                with tempfile.NamedTemporaryFile(dir="/dev/shm", delete=False, suffix=".py") as f:
                    f.write(expression.encode("utf-8"))
                    temporary_filename = f.name

                # execute the code
                # collect the output
                # delete the temporary file
                try:
                    value = check_output(["python", temporary_filename], stderr=STDOUT, timeout=0.1)
                    value = eval(value)
                except:
                    value = None
                os.remove(temporary_filename)
            else:
                # quasi sanitation
                # make sure that there is no `import` statement
                # make sure that there is no `open` statement
                # make sure that there is no `exec` statement
                # make sure that there is no `eval` statement
                # make sure that there is no `os` statement
                # make sure that there is no `sys` statement
                # make sure that there is no `subprocess` statement
                has_bad_thing = False
                for bad_thing in ["import", "open(", "exec(", "eval(", "os.", "sys.", "subprocess."]:
                    if bad_thing in expression:
                        has_bad_thing = True
                        print(f"bad thing found in generated code, skipping: `{bad_thing}`")
                        break
                if has_bad_thing:
                    value = None
                
                else:
                    modified_expression = expression.replace("print(", "OUTPUT.append(")
                    
                    try:
                        OUTPUT = func_timeout.func_timeout(0.1, exec_and_return_output, args=(modified_expression,))
                        # we expect exactly one output
                        value = eval(str(OUTPUT[-1]))
                        print("successfully executed")
                    except:
                        print("failed to execute")
                        value = None
        else:
            try:
                value = check_output(["python", "-c", expression], stderr=STDOUT, timeout=0.1)
                value = eval(value)
            except:
                value = None
        
        result = [value]
        buffered_evaluations[expression] = result

    if len(buffered_evaluations) > 1000:
        print("dumping buffered evaluations to disk...")
        with FileLock(f'{fn}.lck'):
            # load the saved evaluations again, in case someone else has updated them
            if os.path.exists(fn):
                with open(fn, "rb") as handle:
                    saved_evaluations = pickle.load(handle)
            
            saved_evaluations.update(buffered_evaluations)
            buffered_evaluations = {}
            with open(fn, "wb") as handle:
                pickle.dump(saved_evaluations, handle)

    return result[0]

cached_eval.py
ellisk42/humanlike_fewshot_learning
"""Beta Feature: base interface for cache."""
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Tuple

from sqlalchemy.sql.expression import func

from sqlalchemy import Column, Integer, String, create_engine, select, ARRAY, Float
from sqlalchemy.engine.base import Engine
from sqlalchemy.orm import Session

try:
    from sqlalchemy.orm import declarative_base
except ImportError:
    from sqlalchemy.ext.declarative import declarative_base



class BaseCache(ABC):
    """Base interface for cache."""

    @abstractmethod
    def lookup(self, prompt: str, llm_string: str):
        """Look up based on prompt and llm_string."""

    @abstractmethod
    def update(self, prompt: str, llm_string: str, return_val):
        """Update cache based on prompt and llm_string."""


class InMemoryCache(BaseCache):
    """Cache that stores things in memory."""

    def __init__(self) -> None:
        """Initialize with empty cache."""
        self._cache = {}

    def lookup(self, prompt: str, llm_string: str):
        """Look up based on prompt and llm_string."""
        return self._cache.get((prompt, llm_string), None)

    def update(self, prompt: str, llm_string: str, return_val) -> None:
        """Update cache based on prompt and llm_string."""
        self._cache[(prompt, llm_string)] = return_val


Base = declarative_base()


class FullLLMCache(Base):  # type: ignore
    """SQLite table for full LLM Cache (all generations)."""

    __tablename__ = "full_llm_cache"
    prompt = Column(String, primary_key=True)
    llm = Column(String, primary_key=True)
    temperature = Column(Float, primary_key=True)
    max_tokens = Column(Integer, primary_key=True)
    stop = Column(String, primary_key=True)
    idx = Column(Integer, primary_key=True)
    response = Column(String)
    log_probability = Column(Float)
    n_tokens = Column(Integer)

class EvaluationCache(Base):  # type: ignore
    """SQLite table for remembering python evals."""

    __tablename__ = "evals"
    expression = Column(String, primary_key=True)
    return_value = Column(String)


class SQLAlchemyCache(BaseCache):
    """Cache that uses SQAlchemy as a backend."""

    def __init__(self, engine: Engine, cache_schema: Any = FullLLMCache):
        """Initialize by creating all tables."""
        self.engine = engine
        self.cache_schema = cache_schema
        self.cache_schema.metadata.create_all(self.engine)

    def lookup(self, prompt: str, llm_string: str, temperature: float, max_tokens: int, stop, n=None):
        """Look up based on prompt and llm_string."""
        if not isinstance(stop, str): stop = "|".join(stop)
        if n is None:
            stmt = (
                select(self.cache_schema.response,
                    self.cache_schema.log_probability,
                    self.cache_schema.n_tokens, 
                    self.cache_schema.idx)
                .where(self.cache_schema.prompt == prompt)
                .where(self.cache_schema.llm == llm_string)
                .where(self.cache_schema.temperature == temperature)
                .where(self.cache_schema.max_tokens == max_tokens)
                .where(self.cache_schema.stop == stop)
                .order_by(self.cache_schema.idx)
            )
        else:
            stmt = (
                select(self.cache_schema.response,
                    self.cache_schema.log_probability,
                    self.cache_schema.n_tokens, 
                    self.cache_schema.idx)
                .where(self.cache_schema.prompt == prompt)
                .where(self.cache_schema.llm == llm_string)
                .where(self.cache_schema.temperature == temperature)
                .where(self.cache_schema.max_tokens == max_tokens)
                .where(self.cache_schema.stop == stop)
                .where(self.cache_schema.idx < n)
                .order_by(self.cache_schema.idx)
                .limit(n)
            )
        with Session(self.engine) as session:
            generations = [row for row in session.execute(stmt)]
            generations.sort(key=lambda x: x[3])
            generations = [(row[0], row[1], row[2]) for row in generations]
            
            if len(generations) > 0:
                return generations
        return None

    def n_entries(self, prompt: str, llm_string: str, temperature: float, max_tokens: int, stop, session=None):
        """Look up based on prompt and llm_string."""
        if not isinstance(stop, str): stop = "|".join(stop)
        stmt = (
            select(func.max(self.cache_schema.idx))
            .where(self.cache_schema.prompt == prompt)
            .where(self.cache_schema.llm == llm_string)
            .where(self.cache_schema.temperature == temperature)
            .where(self.cache_schema.max_tokens == max_tokens)
            .where(self.cache_schema.stop == stop)
        )
        if session is None:
            with Session(self.engine) as session:
                generations = list(session.execute(stmt))
                if len(generations) > 0:
                    assert len(generations) == 1
                    g = generations[0][0]
                    if g is None: return 0
                    return g+1
        else:
            generations = list(session.execute(stmt))
            if len(generations) > 0:
                assert len(generations) == 1
                g = generations[0][0]
                if g is None: return 0
                return g+1
                    
        return 0

    def update(self, prompt: str, llm_string: str, return_val, temperature: float, max_tokens: int, stop) -> None:
        if not isinstance(stop, str): stop = "|".join(stop)
        for i, (generation,ll,n_tokens) in enumerate(return_val):
            item = self.cache_schema(
                prompt=prompt, llm=llm_string, response=generation, idx=i,
                log_probability=ll, n_tokens=n_tokens, 
                temperature=temperature, max_tokens=max_tokens, stop=stop
            )
            with Session(self.engine) as session, session.begin():
                session.merge(item)

    def extend(self, prompt: str, llm_string: str, return_val, temperature: float, max_tokens: int, stop, session=None) -> None:
        if not isinstance(stop, str): stop = "|".join(stop)
        n = self.n_entries(prompt, llm_string, temperature, max_tokens, stop, session=session)
        print(n, "+", len(return_val), "=>", n+len(return_val), "entries")

        new_items = []
        for i, (generation,ll,n_tokens) in enumerate(return_val):
            item = self.cache_schema(
                prompt=prompt, llm=llm_string, response=generation, idx=i+n,
                log_probability=ll, n_tokens=n_tokens, 
                temperature=temperature, max_tokens=max_tokens, stop=stop
            )
            new_items.append(item)

        if session is None:
            with Session(self.engine) as session, session.begin():
                session.add_all(new_items)
                session.commit()
        else:
            print("about to add new items")
            session.add_all(new_items)
            print("about to commit new items")
            session.commit()
            print("committed new items")


class SQLiteCache(SQLAlchemyCache):
    """Cache that uses SQLite as a backend."""

    def __init__(self, database_path: str = ".langchain.db"):
        """Initialize by creating the engine and all tables."""
        engine = create_engine(f"sqlite:///{database_path}")
        super().__init__(engine)

if __name__ == '__main__':
    assert False, "already ran this"
    import pickle
    x=SQLiteCache("completions.db")
    for llm in ["gpt-3.5-turbo", "code-cushman-001", "code-davinci-002",
                "gpt-4", "text-davinci-003"]:
        filepath = f"{llm}-completions.pickle"
        print(f"About to load {filepath}...")
        with open(filepath, "rb") as f:
            completions = pickle.load(f)
        print(f"Loaded {filepath}.")
        print(f"About to insert {sum(len(cs) for cs in completions.values())} entries into database...")
        #with Session(x.engine) as session, session.begin():
        if True:
            for (prompt, temperature, stop, max_tokens), completion in completions.items():
                x.extend(prompt, llm, completion, temperature, max_tokens, stop, session=None)
        print(f"Inserted {len(completions)} entries into database.")

db.py
ellisk42/humanlike_fewshot_learning
import pandas as pd
import os 
import numpy as np
import pickle

probability_in_concept={}

def load_number_humans():
    global probability_in_concept

    for root, dirs, files in os.walk("tenenbaum_data"):
        for file in files:
            if file.endswith(".csv"):
                with open(os.path.join(root, file), "r") as f:
                    examples = tuple(sorted(map(int, file.replace(".csv", "").replace("[", "").replace("]", "").split(","))))
                    
                    probability_in_concept[examples] = {int(new_number): float(in_concept) for line in f for new_number, in_concept in [line.strip().split(",")] }


def get_human_number_data(examples):
    if len(probability_in_concept) == 0:
        load_number_humans()
    canonical = tuple(sorted(examples))
    return [probability_in_concept[canonical].get(j, -1)
            for j in range(100+1) ]

def special_concept(concept):
    # load the human responses from special_human_data.tsv
    # this is a tab-separated file with the following columns:
    # base_64_encoding_of_responses, nl_response, concept
    # the base_64_encoding_of_responses is a base 64 encoding of a list of lists
    human_responses = {}
    with open("special_human_data.tsv") as handle:
        import base64
        lines = handle.readlines()
        for line in lines:
            encoded_responses, nl_response, this_concept = line.strip().split("\t")
            this_concept = int(this_concept)
            decoded_responses = eval(base64.b64decode(encoded_responses).decode("utf-8"))
            if this_concept not in human_responses:
                human_responses[this_concept] = [ [ [r] for r in batch ] for batch in decoded_responses]
            else:
                for i, batch in enumerate(decoded_responses):
                    for j in range(len(batch)):
                        human_responses[this_concept][i][j].append(batch[j])

    # compute average human responses
    human_responses = {concept: [ [np.mean(rs) for rs in batch] for batch in batches] for concept, batches in human_responses.items()}


    import random
    random.seed(69420)

    if concept == 200:
        # most popular color
        examples = []
        learning_curve = []
        for i in range(15):
            dominant_color = ["yellow", "green", "blue"][random.randint(0,2) if i>2 else i]
            num_dominant_objects = random.choice([3,4]) if random.random()>0.3 or i == 0 else random.choice([2,3,4])
            num_minority_objects = random.choice(list(range(1, min(3, num_dominant_objects))))

            minority_color = random.choice([c for c in ["yellow", "green", "blue"] if c != dominant_color])

            shapes_that_have_the_dominant_color = [(shape, dominant_color, sz) for shape in ["triangle", "rectangle", "circle"] for sz in [1, 2, 3]]
            shapes_that_have_the_minority_color = [(shape, minority_color, sz) for shape in ["triangle", "rectangle", "circle"] for sz in [1, 2, 3]]

            positive_examples = random.sample(shapes_that_have_the_dominant_color, num_dominant_objects)
            negative_examples = random.sample(shapes_that_have_the_minority_color, num_minority_objects)

            labels = [True]*num_dominant_objects + [False]*num_minority_objects

            examples.append( list(zip(labels, positive_examples+negative_examples)))
            random.shuffle(examples[-1])

            labels = np.array([l for l, _ in examples[-1]])

            hr = np.array(human_responses[200][i])

            learning_curve.append(hr * np.array(labels) + (1-np.array(labels)) * (1-hr) )
        return learning_curve, "(lambda (S) (= (color x)) (most-popular-color S))", examples

    if concept == 202:
        # most popular shape
        examples = []
        learning_curve = []
        import random
        for i in range(15):
            dominant_shape = ["triangle", "rectangle", "circle"][random.randint(0,2) if i>2 else i]
            num_dominant_objects = random.choice([3,4]) if random.random()>0.3 or i ==0 else random.choice([2,3,4])
            num_minority_objects = random.choice(list(range(1, min(3, num_dominant_objects))))
            minority_shape = random.choice([c for c in ["triangle", "rectangle", "circle"] if c != dominant_shape])

            shapes_that_have_the_dominant_shape = [(dominant_shape, color, sz) for color in ["yellow", "green", "blue"] for sz in [1, 2, 3]]
            shapes_that_have_the_minority_shape = [(minority_shape, color, sz) for color in ["yellow", "green", "blue"] for sz in [1, 2, 3]]

            positive_examples = random.sample(shapes_that_have_the_dominant_shape, num_dominant_objects)
            negative_examples = random.sample(shapes_that_have_the_minority_shape, num_minority_objects)

            labels = [True]*num_dominant_objects + [False]*num_minority_objects

            examples.append( list(zip(labels, positive_examples+negative_examples)))
            random.shuffle(examples[-1])

            labels = np.array([l for l, _ in examples[-1]])

            hr = np.array(human_responses[202][i])

            learning_curve.append(hr * np.array(labels) + (1-np.array(labels)) * (1-hr) )

        return learning_curve, "(lambda (S) (= (shape x)) (most-popular-shape S))", examples

    # now we do the least popular color and shape
    # these are the same as the most popular color and shape, but with the labels flipped
    if concept == 201:
        # least popular color
        examples = []
        learning_curve = []
        for i in range(15):
            minority_color = ["yellow", "green", "blue"][random.randint(0,2) if i>2 else i]
            num_dominant_objects = random.choice([3,4]) if random.random()>0.3 or i == 0 else random.choice([2,3,4])
            num_minority_objects = random.choice(list(range(1, min(3, num_dominant_objects))))
            dominant_color = random.choice([c for c in ["yellow", "green", "blue"] if c != minority_color])

            shapes_that_have_the_dominant_color = [(shape, dominant_color, sz) for shape in ["triangle", "rectangle", "circle"] for sz in [1, 2, 3]]
            shapes_that_have_the_minority_color = [(shape, minority_color, sz) for shape in ["triangle", "rectangle", "circle"] for sz in [1, 2, 3]]

            positive_examples = random.sample(shapes_that_have_the_dominant_color, num_dominant_objects)
            negative_examples = random.sample(shapes_that_have_the_minority_color, num_minority_objects)

            labels = [False]*num_dominant_objects + [True]*num_minority_objects

            examples.append( list(zip(labels, positive_examples+negative_examples)))

            random.shuffle(examples[-1])

            labels = np.array([l for l, _ in examples[-1]])

            hr = np.array(human_responses[201][i])

            learning_curve.append(hr * np.array(labels) + (1-np.array(labels)) * (1-hr) )

        return learning_curve, "(lambda (S) (= (color x)) (least-popular-color S))", examples

    if concept == 203:
        # least popular shape
        examples = []
        learning_curve = []
        for i in range(15):
            minority_shape = ["triangle", "rectangle", "circle"][random.randint(0,2) if i>2 else i]
            num_dominant_objects = random.choice([3,4]) if random.random()>0.3 or i == 0 else random.choice([2,3,4])
            num_minority_objects = random.choice(list(range(1, min(3, num_dominant_objects))))
            dominant_shape = random.choice([c for c in ["triangle", "rectangle", "circle"] if c != minority_shape])

            shapes_that_have_the_dominant_shape = [(dominant_shape, color, sz) for color in ["yellow", "green", "blue"] for sz in [1, 2, 3]]
            shapes_that_have_the_minority_shape = [(minority_shape, color, sz) for color in ["yellow", "green", "blue"] for sz in [1, 2, 3]]

            positive_examples = random.sample(shapes_that_have_the_dominant_shape, num_dominant_objects)
            negative_examples = random.sample(shapes_that_have_the_minority_shape, num_minority_objects)

            labels = [False]*num_dominant_objects + [True]*num_minority_objects

            examples.append( list(zip(labels, positive_examples+negative_examples)))
            random.shuffle(examples[-1])

            labels = np.array([l for l, _ in examples[-1]])

            hr = np.array(human_responses[203][i])

            learning_curve.append(hr * np.array(labels) + (1-np.array(labels)) * (1-hr) )

        return learning_curve, "(lambda (S) (= (shape x)) (least-popular-shape S))", examples



def get_learning_curve(concept, ordering):
    
    if concept >= 200:
        # these are special new concepts that Kevin added
        return special_concept(concept)

    os.system(f"mkdir -p set_data/precomputed_pickles")
    pickle_filename = f"set_data/precomputed_pickles/hg{concept:02}_L{ordering}.pickle"
    if os.path.exists(pickle_filename):
        with open(pickle_filename, "rb") as handle:
            return pickle.load(handle)       

    print("get_learning_curve", concept, ordering, "being called for the first time. sorry this is going to take a while but next time it will be cached") 

    data = pd.read_csv("set_data/TurkData-Accuracy.txt", sep="\t", index_col=False, low_memory=False)

    data = data[data["concept"] == f"hg{concept:02}"]
    
    data = data[data["list"] == f"L{ordering}"]

    curve = []

    n_sets = data["set.number"].max()

    subjects = data["subject"].unique()  

    # for each subject, get the average accuracy
    subject_accuracy = {}
    for subject in subjects:
        this_subject = data[data["subject"] == subject]
        subject_accuracy[subject] = (this_subject["response"] == this_subject["right.answer"]).mean()
    
    # remove the subjects with low accuracy, defined as <2 standard deviations below the mean
    import numpy as np
    mean = np.mean(list(subject_accuracy.values()))
    std = np.std(list(subject_accuracy.values()))
    
    bad_accuracies = [accuracy for subject, accuracy in subject_accuracy.items() if accuracy <= mean - 2*std]
    bad_subjects = [subject for subject, accuracy in subject_accuracy.items() if accuracy <= mean - 2*std]

    # remove the subjects with low accuracy, defined as the bottom third in terms of accuracy
    # bad_subjects = [subject for subject, accuracy in subject_accuracy.items() if accuracy <= np.percentile(list(subject_accuracy.values()), 33)]

    # also remove any subjects who completed fewer than 5 sets
    for subject in subjects:
        this_subject = data[data["subject"] == subject]
        if len(this_subject["set.number"].unique()) < 5:
            bad_subjects.append(subject)

    for set_number in range(1, n_sets+1):
        this_set = data[data["set.number"] == set_number]
        this_set = this_set[~this_set["subject"].isin(bad_subjects)]
        # n_responses = this_set["response.number"].max()
        # for response in range(1, n_responses+1):
        #     this_response = this_set[this_set["response.number"] == response]
        
        this_response = this_set

        n_probes = this_response["response.number"].max()
                
        probe_accuracy=[]
        for probe in range(1, n_probes+1):
            this_probe = this_response[this_response["response.number"] == probe]
            probe_accuracy.append((this_probe["response"] == this_probe["right.answer"]).mean())
        # if True:
        #     accuracy = (this_response["response"] == this_response["right.answer"]).mean()
        
        curve.append(probe_accuracy)

    # and now we load the actual examples, and the concept
    with open(f"set_data/concepts/CONCEPT_hg{concept:02}__LIST_L{ordering}.txt") as handle:
        lines = handle.readlines()
    
    expression = lines[0].strip()
    examples = []
    for line_number, line in enumerate(lines[1:]):
        line = line.strip().split("\t")
        truth_values = [c == "t" for c in line[0] if c in "tf" ]
        objects = [ (features[0], features[1], int(features[2]))
                    for object_features in line[1:]
                    for features in [object_features.split(",")]]
        assert len(truth_values) == len(objects)
        assert line_number >= len(curve) or len(truth_values) == len(curve[line_number])
        examples.append(list(zip(truth_values, objects)))
    
    expression = expression.replace("*", "").replace("eqv", "=").replace("'", "").replace("(lambda (x) ", "")[:-1]

    with open(pickle_filename, "wb") as handle:
        pickle.dump((curve, expression, examples), handle)

    return curve, expression, examples

def easy_concepts(N):
    # compute top N easy higher order concepts, according to humans    
    averages = {}
    for n in range(1, 112+1):
        c1, expression, _ = get_learning_curve(n, 1)
        data = c1+get_learning_curve(n, 2)[0]
        # flatten data which is a list of lists
        data = [d for d_ in data for d in d_]
        averages[n] = np.mean(data)
    #sort them by average and print them in descending order
    print(list(sorted(averages.items(), key=lambda x: x[1], reverse=True)))
    easy = []
    for n, average in sorted(averages.items(), key=lambda x: x[1], reverse=True):
        easy.append(n)
        if len(easy) == N: break
    
    return easy

if __name__ == '__main__':
    # dump out the special concepts in this format
    # hg24 L1 10 1 False rectangle blue 2 1 21
    # hg24 L1 11 1 False circle yellow 3 0 22
    # hg24 L1 11 2 False rectangle blue 3 0 22
    # hg24 L1 12 1 False rectangle yellow 1 0 22
    for ll in [1,2]:
        for n in [200,201]:
            human_learning_curve, _, examples  = special_concept(n)
            
            for i, (ex, human_accuracies) in enumerate(zip(examples, human_learning_curve)):
                human_accuracies *= 14 # I know there were fourteen subjects

                for j in range(len(ex)):
                    flag, (shape, color, sz) = ex[j]
                    a = human_accuracies[j]
                    yes = int(a if flag else 14-a)
                    no = int(14-yes)
                    print(f"hg{n:02} L{ll} {i+1} {j+1} {flag} {shape} {color} {sz} {yes} {no}")

    assert False
    special_concepts = {}
    for n in [200,201,202,203]:
        examples = special_concept(n)[-1]
        examples=[ [ [flag, {"color":color, "shape":shape, "size":["small","medium","large"][sz-1]}] for flag, (shape, color, sz) in ex] for ex in examples]
        special_concepts[n]=examples
        print(n, examples)
    serialization = str(special_concepts).replace("True", "true").replace("False", "false").replace("'", '"')
    # dump to human_experiment_webpage/data.json
    with open("human_experiment_webpage/data.json", "w") as handle: 
        handle.write(serialization)
    print(serialization)
    assert False
    # load all of these set concepts
    import numpy as np

    averages = {}
    for n in range(26, 112+1):
        data = get_learning_curve(n, 1)[0]+get_learning_curve(n, 2)[0]
        # flatten data which is a list of lists
        data = [d for d_ in data for d in d_]
        print(f"Concept {n}: {np.mean(data)}")
        averages[n] = np.mean(data)
    # sort them by average and print them in descending order
    for n, average in sorted(averages.items(), key=lambda x: x[1], reverse=True):
        print(f"Concept {n}: {average}")
    assert False


    data = get_human_number_data( [30, 31, 33, 24, 21, 36, 39])
    import pdb; pdb.set_trace()

    y, expression, examples = get_learning_curve(19,1)
    print([len(ex) for ex in examples])
    print([len(_y) for _y in y])
    assert False
    import matplotlib.pyplot as plt
    plt.figure()
    
    # flatten y 
    y = [y__ for y_ in y for y__ in y_]
    plt.plot(range(len(y)), y)
    for j, ex in enumerate(examples[:10]):
        print(f"Example set {j}:")
        for flag, features in ex:
            print(features, "\t", flag)
        print()


    plt.title(expression)
    # set the y-limits to the range 0-1
    plt.ylim(0,1)
    plt.show()

    from proposal import propose_set_hypothesis
    for rule in propose_set_hypothesis(examples[:100], 1000, temperature=1):
        print(rule[0])

    print(expression)

human.py
ellisk42/humanlike_fewshot_learning
import os
import pickle
from subprocess import STDOUT, check_output
import math

from proposal import sample_completions, shape_examples_prompt, higher_order_examples_text
from utilities import logsumexp

def execution_likelihood(h, evidence, epsilon=None):

    support = compute_support(h)

    if not isinstance(evidence, list):
        evidence = [evidence]

    L = 0.
    for n in evidence:
        if n in support:
            if epsilon is None:
                L -= math.log(len(support))
            else:
                L += logsumexp([math.log(epsilon) - math.log(100), 
                                math.log(1-epsilon) - math.log(len(support))])
        else:
            if epsilon is None:
                return float("-inf")
            else:
                L += (math.log(epsilon) - math.log(100))
            
    return L

def save_likelihood_cache():
    global SUPPORT_CACHE
    if SUPPORT_CACHE is None: return 
    with open("support.pickle", "wb") as handle:
        pickle.dump(SUPPORT_CACHE, handle)

        
SUPPORT_CACHE=None
        
def compute_support(h):
    global SUPPORT_CACHE

    if SUPPORT_CACHE is None:
        SUPPORT_CACHE={}
        if os.path.exists("support.pickle"):
            with open("support.pickle", "rb") as handle:
                SUPPORT_CACHE = pickle.load(handle)
            print("Loaded", len(SUPPORT_CACHE), "support calculations")
        else:
            print("cannot find cashed support calculations")

    if h in SUPPORT_CACHE: return SUPPORT_CACHE[h]

    if "#" in h:
        h = h[:h.index("#")]
    try:
        support = check_output(["python", "-c", f"import math; print([{h} for num in range(100+1) ])"], stderr=STDOUT, timeout=0.1)
        support = eval(support)
        support = [n for n, flagged in enumerate(support) if flagged]
    except:
        support = []

    SUPPORT_CACHE[h] = support

    return support



def nl_likelihood(h, n):
    """for each number 1..100, returns a probability [0,1]"""

    if isinstance(n, list):
        total_likelihood=1
        for x in n:
            total_likelihood*=nl_likelihood(h, x)
        return total_likelihood
    
    assert isinstance(n, int)
    
    completions = sample_completions(f"""Question: Regarding the number {n}, is it {h}?
Answer (one word, yes/no): """, 30, max_tokens=2, engine="gpt-3.5-turbo",)
    
    yeses = sum( text.startswith("y") or text.startswith("Y")
                 for text, _, _ in completions )
    nos = sum( text.startswith("n") or text.startswith("N")
                 for text, _, _ in completions )

    if yeses == 0 or (yeses == 0 and nos == 0):
        value=float("-inf")
    else:
        value=math.log(yeses/(nos+yeses))

    return value
        


def marginal_lm_likelihood(examples, probe, n=30, temperature=1, domain="number"):
    if domain == "number":

        # now we do this with gp4
        completions = sample_completions(
            engine="gpt-4",
            prompt=f"""Here are a few example number concepts:
-- The number is even
-- The number is between 30 and 45
-- The number is a power of 3
-- The number is less than 10

Here are some random examples of numbers belonging to a possibly different number concept:
{', '.join(map(str,examples))}

Question: Does the number {probe} belong to the same concept as the above numbers?
Answer (one word, yes/no):""",
            temperature=temperature,   
            n=10,
            stop="\n",
            max_tokens=2)

#         completions = sample_completions(
#         engine="code-cushman-001",
#         prompt=f"""# Python 3
# # Here are a few example number concepts:
# # -- The number is even
# # -- The number is between 30 and 45
# # -- The number is a power of 3
# # -- The number is less than 10
# # 
# # Here are some random examples of numbers belonging to a different number concept:
# # {', '.join(map(str,examples))}
# # Question: Does the number {probe} belong to the same concept as the above numbers?
# # Answer: """,
#         temperature=1,
#         n=30,
#         stop="\n",
#         max_tokens=2
#         )
    elif domain == "shapes":
        def verbalize_features(f):
            shape, color, size = f
            size = ["small", "medium", "large"][size-1]
            return f"{size} {color} {shape}"
        completions = []
        for examples in shape_examples_prompt(examples):
            prompt = f"""Here are some example concepts defined by a logical rule:

Rule: a triangle.
Rule: a green rectangle.
Rule: big or a rectangle (unless that rectangle is blue).
Rule: not both big and green.
Rule: either big or green, but not both.
Rule: either a rectangle or not yellow.
Rule: a circle.

Now please look at the following examples for a new logical rule.

{examples}

Question: Based on the above examples, is a {verbalize_features(probe)} in the concept?
Answer (one word, just write yes/no):"""

            completions.extend(sample_completions(prompt=prompt,
                                    n=5,
                                    engine="gpt-4", #, #"code-davinci-002",#"gpt-3-turbo",
                                    temperature=temperature,
                                    stop="\n",
                                    max_tokens=2))
    elif domain == "higher_order_shapes":
        def verbalize_features(f):
            shape, color, size = f
            size = ["small", "medium", "large"][size-1]
            return f"({size} {color} {shape})"

        all_the_probes = " ".join(map(verbalize_features, probe))
        probe_probabilities = []

        examples = higher_order_examples_text(examples)
        for this_probe in probe:
            prompt = f"""Here are some example concepts defined by a logical rule:

Rule for Concept #1: Something is positive if it is the biggest yellow object in the example
Rule for Concept #2: Something is positive if there is another object with the same color in the example
Rule for Concept #3: Something is positive if it is the same color as the smallest triangle in the example

Now please look at the following examples for a new logical rule.

{examples}

Now we get a new collection of examples for Concept #4:
{all_the_probes}
Question: Based on the above example, is a {verbalize_features(this_probe)} in the concept?
Answer (one word, just write yes/no):"""

            print(prompt)

            completions = sample_completions(prompt=prompt,
                                    n=10,
                                    engine="gpt-4", #, #"code-davinci-002",#"gpt-3-turbo",
                                    temperature=1, 
                                    stop="\n",
                                    max_tokens=2)
            yeses = sum( text.lower().strip().startswith("y") for text, _, _ in completions )
            nos = sum( text.lower().strip().startswith("n") for text, _, _ in completions )
            if yeses == 0 or (yeses == 0 and nos == 0):
                probe_probabilities.append(float("-inf"))
            else:
                probe_probabilities.append(math.log(yeses/(nos+yeses)))
        return probe_probabilities

    yeses = sum( text.strip().startswith("y") or text.strip().startswith("Y")
                 for text, _, _ in completions )
    nos = sum( text.strip().startswith("n") or text.strip().startswith("N")
                 for text, _, _ in completions )

    if yeses == 0 or (yeses == 0 and nos == 0):
        return float("-inf")

    #print(f"P_gpt4({probe} | {examples}) = {yeses/(nos+yeses)})")    

    return math.log(yeses/(nos+yeses))

def propose_next_number(examples, n, temperature=1):
    # https://dspace.mit.edu/bitstream/handle/1721.1/16714/42471842-MIT.pdf?sequence=2&isAllowed=y
    completions = sample_completions(
        engine="code-davinci-002",
        prompt=f"""# Python 3
# Here are a few example number concepts:
# -- The number is even
# -- The number is between 30 and 45
# -- The number is a power of 3
# -- The number is less than 10
# 
# Here are some random examples of numbers belonging to a different number concept:
# {', '.join(map(str,examples))},""",
        temperature=1,
        n=n,
        stop=",",
        max_tokens=5
    )
    histogram = {}
    total=0
    for c,_,_ in completions:
        try:
            v = int(c.strip().split()[0])
        except: continue
        histogram[v] = 1+histogram.get(v,0)
        total+=1

    return [ math.log(histogram[v]/total) if v in histogram else float("-inf")
             for v in range(101) ]
    


def transpiler_likelihood(nl, evidence, epsilon=None):
    completion = sample_completions(
        engine="code-davinci-002",
        prompt=f"""# Write a python function to check if a number is {nl}.
def check_number(num):
    return""",
        temperature=0,
        n=1,
        stop="\n",
        max_tokens=128
    )[0][0]
    
    if isinstance(evidence, list):
        return execution_likelihood(completion, evidence, epsilon=epsilon)
    else:
        return execution_likelihood(completion, [evidence], epsilon=epsilon)
    


    

likelihood.py
ellisk42/humanlike_fewshot_learning
"""https://dspace.mit.edu/bitstream/handle/1721.1/16714/42471842-MIT.pdf?sequence=2&isAllowed=y"""
import math
import matplotlib.pyplot as plt
import numpy as np
import pickle
import os
import scipy.special
import random
import scipy.special

from scipy.special import log_softmax

import pytorch_lightning as pl
from pytorch_lightning.loggers import TensorBoardLogger


# import matplotlib
# matplotlib.use('QtAgg')

from proposal import set_sampling_seed
from human import get_human_number_data
from prior import code_prior, nl_prior
from proposal import propose_code_hypotheses, propose_nl_hypotheses, sample_completions
from likelihood import execution_likelihood, save_likelihood_cache, compute_support, nl_likelihood, marginal_lm_likelihood, transpiler_likelihood, propose_next_number
from utilities import logsumexp
from parameter_fitting import NumberModel, number_data_loaders

from collections import namedtuple

Config = namedtuple("Config", "n_proposals epsilon nl_hypotheses code_execution deduplicate reweigh use_prior per_token use_likelihood no_latent propose_from_prior MAP")

def nl2py(nl):
    # compute the engine, defaulting to "code-davinci-002" unless the command line options say "llama"
    engine = "code-davinci-002" if not arguments.llama else "llama"
    return sample_completions(
        engine=engine,
        prompt=f"""# Write a python function to check if a number is {nl}.
def check_number(num):
    return""",
        temperature=0,
        n=1,
        stop="\n",
        max_tokens=128
    )[0][0]

def likelihood(h, E, config):
    if not config.code_execution:
        return nl_likelihood(h, E)

    if config.nl_hypotheses:
        return transpiler_likelihood(h, E, epsilon=config.epsilon)
    else:
        return execution_likelihood(h, E, epsilon=config.epsilon)

def likelihood_in_concept(h, E, config):
    if not config.code_execution:
        return nl_likelihood(h, E)

    if config.nl_hypotheses:
        log_probability = transpiler_likelihood(h, E)
    else:
        log_probability = execution_likelihood(h, E)

    if math.isinf(log_probability):
        return log_probability
    else:
        return 0.

def prior(h, config):
    if config.nl_hypotheses:
        return nl_prior(h)
    else:
        return code_prior(h)

Result = namedtuple("Result", "hypotheses log_prior log_weight support log_posterior predictions")
def important_bayesian_stuff(config, E):

    # in the event that we are doing a non-Bayes model, 
    # we can still pretend it is by saying there has a single latent hypothesis that directly predicts the data
    if config.no_latent:
        probability_in_concept = [0 for _ in range(101) ]
        if config.use_likelihood:
            samples = propose_next_number(E, config.n_proposals)
            for i in range(len(probability_in_concept)):
                probability_in_concept[i] = samples[i]
        else:
            for i in range(len(probability_in_concept)):
                probability_in_concept[i] = marginal_lm_likelihood(E, i)
        probability_in_concept = np.exp(np.array([probability_in_concept]))
        result = Result([("no latent", 0, 0)], np.array([0]), np.array([0]),
                        probability_in_concept, np.array([0]), 
                        probability_in_concept)
        return result

    if config.nl_hypotheses:
        hs = propose_nl_hypotheses(E if not config.propose_from_prior else [], config.n_proposals, 
                                   engine="code-davinci-002")# if not arguments.llama else "llama")
    else:
        assert not config.propose_from_prior
        hs = propose_code_hypotheses(E, config.n_proposals)

    if config.deduplicate:
        hs = {hypothesis: stats for hypothesis, *stats in hs }
        hs = [(hypothesis, *stats) for hypothesis, stats in hs.items()]
    
    print("about to compute priors")
    priors = prior([h for h, q_h, l_h in hs ], config)
    print("done computing priors")

    if config.per_token:
        priors = [ pr/ln for pr, ln in priors]
    else:
        priors = [ pr for pr, ln in priors]

    if config.reweigh:
        log_weight = np.array([ -q_h for h, q_h, l_h in hs ])
    else:
        log_weight = np.array([ 0. for h, q_h, l_h in hs ])

    if config.nl_hypotheses and config.code_execution:
        print("computed priors, about to compute program translations")
        code_hypotheses = {nl: nl2py(nl) for nl in set(nl_ for nl_, _, _ in hs)}
        # convert to list
        code_hypotheses = [code_hypotheses[nl] for nl, _, _ in hs]
    if not config.nl_hypotheses and config.code_execution:
        code_hypotheses = [h for h, q_h, l_h in hs]

    # now we compute support
    print("about to compute support")
    if config.code_execution:
        support = [ [ (n in s)*1 for n in range(101) ] for h in code_hypotheses for s in [compute_support(h)] ]
    else:
        support = [ [ math.exp(nl_likelihood(h, n)) for n in range(101) ]
                     for h, q_h, l_h in hs ]

    save_likelihood_cache()

    print("done computing support")

    support = np.array(support)
    pN_H = support/np.maximum(np.sum(support, -1)[:,None], 1e-10)

    priors = np.array(priors)
    evidence = np.array([n in E for n in range(101) ])

    # likelihood over data includes noise model
    like = np.log(pN_H * (1-config.epsilon) + 0.01 * config.epsilon)
    like = np.sum(like*evidence, -1)
    
    # posterior over hypotheses
    posterior = like + priors + log_weight
    # normalize posterior
    posterior = log_softmax(posterior)

    # construct posterior predictive
    # this is a number between 0 and 1, for each index 0--100
    posterior_predictive = np.sum(support * np.exp(posterior)[:,None], 0)

    if config.nl_hypotheses and config.code_execution:
        hypotheses = list(zip([h for h, q_h, l_h in hs], code_hypotheses))
    else:
        hypotheses = [h for h, q_h, l_h in hs]

    best = np.max(posterior)
    # find top k=3 and show them
    top_indices = np.argsort(posterior)[::-1][:3]
    for i in top_indices[::-1]:
        h = hypotheses[i]
        print("BEST", h)
        print(posterior[i], f"p({E}|h)={like[i]}",
        priors[i], log_weight[i])
        print()
    print()
    #assert False



    return Result(hypotheses, priors, log_weight, support, posterior, posterior_predictive)


def optimized_posterior_predictive(Es, config, prior, neural_network_options, experiment_name):

    results = [important_bayesian_stuff(config, E) for E in Es]

    human_predictions = [get_human_number_data(E) for E in Es]
    

    # we are going to return the average predictions over every fold
    # but never average in the prediction for a human judgment when we trained on that judgment
    # to do that, we are going to collect the predictions for each fold, for each dataset, for each number
    # finally we return the average
    all_predictions = [ [ list() for _ in range(101) ] for _ in Es]

    folds = 10
    for fold, (train, validation) in enumerate(number_data_loaders(results, human_predictions, batch_size=1, folds=folds)):
        model = NumberModel(results, Es, prior=prior, MAP=config.MAP, **neural_network_options)
        logger = TensorBoardLogger("lightning_logs", name=experiment_name+f"_fold_{fold}of{folds}")
        trainer = pl.Trainer(max_epochs=model.iterations, val_check_interval=len(train)*20, check_val_every_n_epoch=None, logger=logger, accelerator="cpu")
        trainer.fit(model, train, validation, ckpt_path=None)

        predictions = [model.posterior_predictive(index).detach().cpu().exp().numpy() for index in range(len(Es))]

        for d in range(len(Es)):
            for n in range(101):
                assert validation.dataset.pairs[d][0] == d

                # we don't have human data for this number, in which case it was never trained on
                # or it was part of validation for this fold
                no_human_data = human_predictions[d][n]<0
                # in theory should be able to get rid of this logic but I don't want to break everything in case it is buggy
                in_validation = validation.dataset.pairs[d][1][n] >= 0
                if no_human_data or in_validation:
                    all_predictions[d][n].append(predictions[d][n])

    # last we train a single model on all the data
    # we do this so that we can get a single posterior over all hypotheses
    # all quantitative metrics are always reported on test data and not on this special model
    model = NumberModel(results, Es, prior=prior, MAP=config.MAP,  **neural_network_options)
    logger = TensorBoardLogger("lightning_logs", name=experiment_name+f"_all_data")
    trainer = pl.Trainer(max_epochs=model.iterations, val_check_interval=None, check_val_every_n_epoch=None, logger=logger, accelerator="cpu")
    trainer.fit(model, number_data_loaders(results, human_predictions, batch_size=1, folds=0), None, ckpt_path=None)

    posteriors = [model.posterior(index).detach().cpu().exp().numpy() for index in range(len(Es))]
    priors = [model.prior(index).detach().cpu().numpy() for index in range(len(Es))]
    likelihoods = [model.likelihood(index).detach().cpu().numpy() for index in range(len(Es))]

    all_posterior_samples = [ list() for _ in range(len(Es)) ]
    for d in range(len(Es)):
        def detuple(hh): return hh[0]  if isinstance(hh, tuple) else hh
        all_posterior_samples[d].extend(list(np.random.choice(list(map(detuple, results[d].hypotheses)), 
                                                                  1000, p=posteriors[d])))
    
    # convert all_posterior_samples to a dictionary mapping a sample to its frequency
    all_posterior_samples = [ {h: all_posterior_samples[d].count(h) for h in set(all_posterior_samples[d])} for d in range(len(Es)) ]

    return [ [ np.mean(np.array(prediction_list)) for prediction_list in predictions ] for predictions in all_predictions ], all_posterior_samples
                    


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description = "")
    parser.add_argument('--n_proposals', type=int, default=100)
    parser.add_argument('--epsilon', "-e", type=float, default=0.1)
    parser.add_argument('--seed', "-s", type=int, default=None)
    parser.add_argument('--export', default=None)
    parser.add_argument('--data', "-d", default="Tenenbaum", choices=["Tenenbaum"])
    parser.add_argument('--per_token', default=False, action="store_true")
    parser.add_argument('--log', default=False, action="store_true")
    parser.add_argument("--prior", default="learned", choices=["learned", "uniform", "fixed"])
    
    parser.add_argument('--hidden', default=0, type=int)
    parser.add_argument('--iterations', default=1000, type=int)
    parser.add_argument('--dropout', default=0.5, type=float)

    parser.add_argument('--propose_from_prior', default=False, action="store_true")
    parser.add_argument('--deduplicate', default=False, action="store_true")
    parser.add_argument('--methods', nargs='+', 
    choices=["latent code", "latent lang", "latent lang2code", 
             "sample+filter", "raw LLM/classify", "raw LLM/generate", "L3", "human"],
    default=["latent code", "latent lang", "latent lang2code", "raw LLM/generate", "raw LLM/classify"])

    parser.add_argument('--llama', default=False, action="store_true")

    arguments = parser.parse_args()

    if arguments.seed is not None:
        set_sampling_seed(arguments.seed)

    # store the neural network options into a dictionary
    neural_network_options = {k:v for k,v in vars(arguments).items() if k in ["hidden", "dropout", "iterations"]}

    nn_options_string ="_"+ "_".join([f"{k}={v}" for k,v in neural_network_options.items()])
    experiment_name = f"numbergame_seed={arguments.seed}_{arguments.data}_n={arguments.n_proposals}{nn_options_string}"
    if arguments.propose_from_prior:
        experiment_name += "_PfP"

    if arguments.llama:
        experiment_name += "_llama"

    configs = {"latent code": Config(n_proposals=arguments.n_proposals,
                                    epsilon=arguments.epsilon,
                                    nl_hypotheses=False,
                                    code_execution=True,
                                    deduplicate=arguments.deduplicate,
                                    reweigh=not arguments.deduplicate,
                                    use_prior=True,
                                    per_token=arguments.per_token,
                                    use_likelihood=True,
                                    no_latent=False,
                                    propose_from_prior=arguments.propose_from_prior,
                                    MAP=False),
               "L3": Config(n_proposals=arguments.n_proposals,
                                    epsilon=arguments.epsilon,
                                    nl_hypotheses=True,
                                    code_execution=True,
                                    deduplicate=arguments.deduplicate,
                                    reweigh=not arguments.deduplicate,
                                    use_prior=False,
                                    per_token=arguments.per_token,
                                    use_likelihood=True,
                                    no_latent=False,
                                    propose_from_prior=arguments.propose_from_prior, 
                                    MAP=True),
              "latent lang": Config(n_proposals=arguments.n_proposals,
                                    epsilon=arguments.epsilon,
                                    nl_hypotheses=True,
                                    code_execution=False,
                                    deduplicate=arguments.deduplicate,
                                    reweigh=not arguments.deduplicate,
                                    use_prior=True,
                                    per_token=arguments.per_token,
                                    use_likelihood=True,
                                    no_latent=False,
                                    propose_from_prior=arguments.propose_from_prior, 
                                    MAP=False),
            "latent lang2code": Config(n_proposals=arguments.n_proposals,
                                    epsilon=arguments.epsilon,
                                    nl_hypotheses=True,
                                    code_execution=True,
                                    deduplicate=arguments.deduplicate,
                                    reweigh=not arguments.deduplicate,
                                    use_prior=True,
                                    per_token=arguments.per_token,
                                    use_likelihood=True,
                                    no_latent=False,
                                    propose_from_prior=arguments.propose_from_prior, 
                                    MAP=False),
            "sample+filter": Config(n_proposals=arguments.n_proposals,
                                    epsilon=arguments.epsilon,
                                    nl_hypotheses=False,
                                    code_execution=True,
                                    deduplicate=False,
                                    reweigh=False,
                                    use_prior=False,
                                    per_token=arguments.per_token,
                                    use_likelihood=False,
                                    no_latent=False,
                                    propose_from_prior=False ,
                                    MAP=False),
            "raw LLM/classify": Config(n_proposals=arguments.n_proposals,
                                    epsilon=arguments.epsilon,
                                    nl_hypotheses=False,
                                    code_execution=False,
                                    deduplicate=False,
                                    reweigh=False,
                                    use_prior=False,
                                    per_token=arguments.per_token,
                                    use_likelihood=False,
                                       no_latent=True,
                                    propose_from_prior=False, 
                                    MAP=False),
               "raw LLM/generate": Config(n_proposals=arguments.n_proposals,
                                    epsilon=arguments.epsilon,
                                    nl_hypotheses=False,
                                    code_execution=False,
                                    deduplicate=False,
                                    reweigh=False,
                                    use_prior=False,
                                    per_token=arguments.per_token,
                                    use_likelihood=True,
                                    no_latent=True,
                                    propose_from_prior=False, 
                                    MAP=False)}
    datasets = [[16, 8, 2, 64],
                [60],
                [16, 23, 19, 20],
                [16],
                [60, 80, 10, 30],
                [60, 52, 57, 55],
                [98,81,86,93],
                [25,4,36,81],
            ]

    # compute the posterior predictive curves for each method on each dataset
    posterior_predictives = {}
    posterior_samples = {}
    for method in arguments.methods:
        if method == "human":
            posterior_predictives.update({(method, tuple(E)): get_human_number_data(E)
                                            for E in datasets})
            continue
        config = configs[method]
        pps, these_posterior_samples = optimized_posterior_predictive(datasets, config, arguments.prior, neural_network_options, experiment_name)
        posterior_predictives.update({(method, tuple(E)): pp
                                       for E, pp in zip(datasets, pps)})
        posterior_samples.update({(method, tuple(E)): this_posterior_samples
                                        for E, this_posterior_samples in zip(datasets, these_posterior_samples)})

    # now we are going to plot everything

    # make the horizontal size of the figure proportional to the number of methods, 
    # and the vertical size proportional to the number of datasets
    plt.figure(figsize=(len(arguments.methods)*8, len(datasets)*2))
    index=1
    for E in datasets:
        for method in arguments.methods:
            pp = posterior_predictives[(method, tuple(E))]
            
            plt.subplot(len(datasets), len(arguments.methods), index)
            index+=1
            for e in E: pp[e] = 1
            plt.bar(list(range(len(pp)))[1:], pp[1:], color ='blue',
                    width = 0.4)
            if 'human' in arguments.methods:
                pp = np.array(posterior_predictives[('human', tuple(E))])
                pp[pp < 0] = 0
                plt.bar(np.array(list(range(len(pp)))[1:])+0.4, pp[1:], color ='green',
                    width = 0.4)
            plt.title(f"{E}\n{method}")
    if arguments.export:
        plt.tight_layout()
        plt.savefig(arguments.export+".png")

        # also save the posterior predictive probabilities in a csv
        with open(arguments.export+".csv", "w") as f:
            f.write("method,dataset,example,pp\n")
            for (method, E), pp in posterior_predictives.items():
                for i, p in enumerate(pp):
                    f.write(f"{method},{'_'.join(map(str, E))},{i},{p}\n")

        # also save the posterior samples in a tsv
        with open(arguments.export+"_samples.tsv", "w") as f:
            f.write("method\tdataset\tsample\tfrequency\n")
            for (method, E), samples in posterior_samples.items():
                for frequency, sample in samples.items():
                    f.write(f"{method}\t{'_'.join(map(str, E))}\t{sample}\t{frequency}\n")
    else:
        plt.show()

    # show a figure with model predictions on the horizontal axis and human predictions on the vertical axis
    
    correlations = {} # one for each method
    for method in arguments.methods:
        if method == "human":
            continue

        relevant_posterior_predictive = {k: v for k, v in posterior_predictives.items() if k[0] == method}
        for (method, E), pp in relevant_posterior_predictive.items():
            human = get_human_number_data(E)
            for n in range(1, 101):
                if n in E: continue
                if human[n] < 0: continue
                correlations[method] = correlations.get(method, list())
                correlations[method].append([pp[n], human[n]])
    plt.figure()
    for method, data in correlations.items():
        data = np.array(data)
        r = np.corrcoef(data[:,0], data[:,1])[0,1]
        plt.scatter(data[:,0], data[:,1], label=method+f"(r={r:.2f})")
        
    plt.plot([0,1], [0,1], color="black")
    plt.xlabel("model prediction")
    plt.ylabel("human prediction")
    plt.legend()
    if arguments.export:
        plt.tight_layout()
        plt.savefig(arguments.export+"_correlations.png")
    else:
        plt.show()

number.py
ellisk42/humanlike_fewshot_learning
from textfeatures import textfeatures

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset

import math
import numpy as np
import random
import pickle


# we are using lightning so import that
import pytorch_lightning as pl

def inverse_sigmoid(z):
    if isinstance(z, torch.Tensor):
        return torch.log(z / (1 - z))
    else:
        return math.log(z / (1 - z))

class NumberModel(pl.LightningModule):
    def __init__(self, precomputed_results, evidence, prior, MAP, hidden=0, dropout=0., iterations=100):
        super(NumberModel, self).__init__()

        assert prior in ["uniform", "fixed", "learned"]
        self.prior_type = prior

        # Create learnable parameters for constants:
        # epsilon, alpha, beta, gamma
        # Each of these is a scalar.
        # epsilon is a probability of a binary event, represented as a binary logit.
        self.epsilon = nn.Parameter(torch.tensor(inverse_sigmoid(0.01)))
        
        self.posterior_temperature = nn.Parameter(torch.tensor(1.))

        self.MAP = MAP
        if self.MAP: assert prior == "uniform"

        self.no_latent = any(result.hypotheses == [("no latent", 0, 0)] for result in precomputed_results )
        
        self.temperature = nn.Parameter(torch.tensor(1.))
        self.offset = nn.Parameter(torch.tensor(0.))

        # Remember definition of results
        # Result = namedtuple("Result", "hypotheses log_prior log_weight support log_posterior predictions")

        self.canonically_number_hypotheses(precomputed_results)

        if self.prior_type == "learned":
            if hidden == 0:
                self.prior_network = nn.Sequential(nn.Dropout(dropout), nn.Linear(384, 1))
            else:
                self.prior_network = nn.Sequential(nn.Dropout(dropout), nn.Linear(384, hidden), nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden, 1))

            print("precomputing text features")
            features = textfeatures([hypothesis[0].lower() if isinstance(hypothesis, tuple) else hypothesis
                                    for hypothesis in self.every_hypothesis])
            features = {hypothesis: features[i] for i, hypothesis in enumerate(self.every_hypothesis)}
            self.precomputed_features = [ torch.tensor([ features[hypothesis]
                                                           for hypothesis in result.hypotheses ])
                                          for result in precomputed_results ]

        self.iterations = iterations

        self.pN_H = [ torch.tensor(result.support/np.maximum(np.sum(result.support, -1)[:,None], 1e-10)) for result in precomputed_results]
        self.evidence = [torch.tensor([n in E for n in range(101) ]) for E in evidence]
        self.precomputed_results = [self.to_torch(result) for result in precomputed_results]

        self.loss_type = "ll"

    def canonically_number_hypotheses(self, precomputed_results):
        # Go through the results, collect all the hypotheses, and deduplicate them.
        # This is going to be used to create a canonical numbering scheme for the hypotheses.
        every_hypothesis = []
        for result in precomputed_results:
            every_hypothesis.extend(result.hypotheses)
        self.every_hypothesis = list(set(every_hypothesis))
        self.hypothesis_numbering_scheme = {hypothesis: i for i, hypothesis in enumerate(self.every_hypothesis)}
        self.precomputed_hypotheses = [ [self.hypothesis_numbering_scheme[h] for h in result.hypotheses]
                                        for result in precomputed_results ]

    def to_torch(self, result):
        return result._replace(**{k: torch.tensor(v).to(self.device)*1 for k,v in result._asdict().items() if isinstance(v, np.ndarray)})

    def prior(self, d):
        # returns log prior for dataset d
        result = self.precomputed_results[d]

        if self.prior_type == "table":
            prior = self.prior_table[self.precomputed_hypotheses[d]] + result.log_prior
        elif self.prior_type == "fixed":
            prior = result.log_prior
        elif self.prior_type == "uniform":
            prior = torch.zeros_like(result.log_prior)
        elif self.prior_type == "learned":
            prior = self.prior_network(self.precomputed_features[d]).squeeze(-1)

        return prior

    def likelihood(self, d):
        # returns log likelihood for dataset d (across all hypotheses)
        result = self.precomputed_results[d]
        E = self.evidence[d]

        # old code: prone to numerical problems because it does not use logsumexp
        # like = torch.log(self.pN_H[d] * F.sigmoid(-self.epsilon) + 0.01 * F.sigmoid(self.epsilon))
        # new code: uses logsumexp to avoid numerical problems
        like = torch.logaddexp(torch.log(self.pN_H[d]) + F.logsigmoid(-self.epsilon), np.log(0.01) + F.logsigmoid(self.epsilon))
        like = torch.sum(like*(E.unsqueeze(0)), -1)

        return like

    def posterior(self, d):
        # returns log posterior for dataset d (across all hypotheses)
        prior_weight, likelihood_weight, proposal_weight = self.posterior_temperature, self.posterior_temperature, self.posterior_temperature
        if self.prior_type == "learned":
            prior_weight = 1 # putting a coefficient on it is redundant
        unnormalized = prior_weight*self.prior(d) + \
                       likelihood_weight*self.likelihood(d) +\
                       proposal_weight*self.precomputed_results[d].log_weight

        if self.MAP:
            return (unnormalized == torch.max(unnormalized, 0)[0]).float().log()
        else:
            return unnormalized-torch.logsumexp(unnormalized, 0)

    def posterior_predictive(self, d):
        if not self.no_latent:
            log_in_concept = self.precomputed_results[d].support.log().clamp(min=-10000) # clamping fixes numerical issues
            posterior_predictive = torch.logsumexp(log_in_concept + self.posterior(d)[:,None], 0)
            #return posterior_predictive
            posterior_predictive = posterior_predictive.exp()
        else:
            # now we have some "no latent" models, aka raw LLMs
            posterior_predictive = self.precomputed_results[d].support.squeeze(0)

        # models without latent variables and get some extra benefits from this transformation
        # only seems fair to give them more parameters to fit

        # with probability (1-epsilon), the model returns the above
        # with probability epsilon, the model returns a uniform distribution
        posterior_predictive = torch.sigmoid(-self.epsilon) * posterior_predictive + \
                                 torch.sigmoid(self.epsilon) * 0.01

        #return posterior_predictive.log()

        # compute the inverse sigmoid of posterior_predictive
        posterior_predictive = torch.logit(posterior_predictive)

        # linear transform
        posterior_predictive = (posterior_predictive - self.offset) / self.temperature

        # logsigmoid
        posterior_predictive = F.logsigmoid(posterior_predictive)

        return posterior_predictive

    def losses(self, d, human_judgments):
        # human_judgments: vector of length 101. Values are negative when there is missing data. Loss is not computed on missing data.
        # returns both the negative log likelihood and mean squared error

        pp = self.posterior_predictive(d)

        # log probability of returning yes/no for each number being in the concept
        pp_yes = pp 
        pp_no = torch.log1p(-pp.exp()+1e-10)
        
        ll = human_judgments*pp_yes + (1-human_judgments)*pp_no
        ll = ll*(human_judgments >= 0)
        ll = torch.sum(ll) / torch.sum(human_judgments >= 0)

        mse = (pp.exp()-human_judgments)**2
        mse = mse*(human_judgments >= 0)
        mse = torch.sum(mse) / torch.sum(human_judgments >= 0)

        return -ll, mse

    # configuration of optimizers
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.001)

    def training_step(self, batch, batch_idx):
        ds, human_accuracies = batch 
        # ds is a list of dataset indices.

        mse_loss, ll_loss = 0, 0
        for d, human_accuracy in zip(ds, human_accuracies):
            ll, mse = self.losses(d, human_accuracy)
            mse_loss += mse
            ll_loss += ll

        if self.loss_type == "mse":
            loss = mse_loss
        elif self.loss_type == "ll":
            loss = ll_loss 

        total_size = len(ds)

        self.log('mse_loss', mse_loss/total_size, on_epoch=True, batch_size=len(ds))
        self.log('ll_loss', ll_loss/total_size, on_epoch=True, batch_size=len(ds))
        self.log('loss', loss/total_size, on_epoch=True, batch_size=len(ds))

        return loss/total_size

    def validation_step(self, batch, batch_idx):
        ds, human_accuracies = batch 
        # ds is a list of dataset indices.

        mse_loss, ll_loss = 0, 0
        human_predictions, machine_predictions = [], []
        for d, human_accuracy in zip(ds, human_accuracies):
            ll, mse = self.losses(d, human_accuracy)
            mse_loss += mse
            ll_loss += ll

            human_accuracy = human_accuracy.detach().cpu().numpy()
            pp = self.posterior_predictive(d).exp().detach().cpu().numpy()
            for i in range(len(pp)):
                if human_accuracy[i] >= 0:
                    human_predictions.append(human_accuracy[i])
                    machine_predictions.append(pp[i])

        if self.loss_type == "mse":
            loss = mse_loss
        elif self.loss_type == "ll":
            loss = ll_loss 

        # compute correlation coefficient
        human_predictions = np.array(human_predictions)
        machine_predictions = np.array(machine_predictions)
        correlation_coefficient = np.corrcoef(human_predictions, machine_predictions)[0,1]

        total_size = len(ds)

        self.log('val_ll_loss', mse_loss/total_size, on_epoch=True, batch_size=len(ds))
        self.log('val_mse_loss', ll_loss/total_size, on_epoch=True, batch_size=len(ds))
        self.log('val_loss', loss/total_size, on_epoch=True, batch_size=len(ds))
        self.log('val_correlation_coefficient', correlation_coefficient, on_epoch=True, batch_size=len(ds))

        return loss/total_size
    



class ShapeModel(pl.LightningModule):
    def __init__(self, precomputed_results, prior, hidden=0, dropout=0., iterations=100, performance=False, load_fixed_prior=False):
        super(ShapeModel, self).__init__()

        assert prior in ["uniform", "fixed", "table", "learned"]
        
        # Create learnable parameters for constants:
        # epsilon_train, epsilon_test, base_rate, memory_decay, likelihood_coefficient, prior_coefficient
        # Each of these is a scalar.
        # Probabilities of binary events are represented as binary logits.
        self.epsilon_train = nn.Parameter(torch.tensor(inverse_sigmoid(0.01)))
        self.epsilon_test = nn.Parameter(torch.tensor(inverse_sigmoid(0.1)))
        self.base_rate = nn.Parameter(torch.tensor(inverse_sigmoid(0.5)))
        self.memory_decay = nn.Parameter(torch.tensor(1.))
        self.likelihood_coefficient = nn.Parameter(torch.tensor(1.0))
        
        self.performance = performance # should we optimize performance instead of fit to humans?

        self.platt_temperature = nn.Parameter(torch.tensor(1.))
        self.platt_offset = nn.Parameter(torch.tensor(0.))

        # torch everything
        self.precomputed_results = [ [ self.to_torch(result) for result in results ]
                                     for results in precomputed_results ]

        self.canonically_number_hypotheses(self.precomputed_results)

        self.prior_type = prior
        if self.prior_type == "table":
            self.prior_table = nn.Parameter(torch.zeros(len(self.hypothesis_numbering_scheme), dtype=torch.float32))
            print("Prior table size:", len(self.hypothesis_numbering_scheme))
        elif self.prior_type == "learned":
            self.load_fixed_prior = load_fixed_prior
            if hidden == 0:
                self.prior_network = nn.Sequential(nn.Dropout(dropout), nn.Linear(384, 1))
                if load_fixed_prior:
                    with open("transferred_prior.pickle", "rb") as handle:
                        transferred_prior_network = pickle.load(handle)
                    self.prior_network = transferred_prior_network
                    for parameter in self.prior_network:
                        parameter.requires_grad_(False)
                    parameter.requires_grad_(False)                    
            else:
                self.prior_network = nn.Sequential(nn.Dropout(dropout), nn.Linear(384, hidden), nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden, 1))

            print("precomputing text features")
            def clean_text(text):
                return text.lower().replace("rule: ", "It is ").replace("something is positive if it", "It is").replace("rule for concept #4: ", "").replace(".", "")
            features = textfeatures([clean_text(hypothesis[0] if isinstance(hypothesis, tuple) else hypothesis)
                                    for hypothesis in self.every_hypothesis])
            features = {hypothesis: features[i] for i, hypothesis in enumerate(self.every_hypothesis)}
            self.precomputed_features = [ [ torch.tensor([ features[hypothesis]
                                                           for hypothesis in result.hypotheses ])
                                            for result in results ]
                                        for results in precomputed_results ]

        

        #self.device = torch.device("cpu") #"cuda:0" if torch.cuda.is_available() else 
        self.loss_type = "ll"
        self.iterations = iterations
        
    def canonically_number_hypotheses(self, precomputed_results):
        # Go through the results, collect all the hypotheses, and deduplicate them.
        # This is going to be used to create a canonical numbering scheme for the hypotheses.
        every_hypothesis = []
        for results in precomputed_results:
            for result in results:
                every_hypothesis.extend(result.hypotheses)
        self.every_hypothesis = list(set(every_hypothesis))
        self.hypothesis_numbering_scheme = {hypothesis: i for i, hypothesis in enumerate(self.every_hypothesis)}
        self.precomputed_hypotheses = [ [ [ self.hypothesis_numbering_scheme[h] for h in result.hypotheses ]
                                            for result in results ]
                                           for results in precomputed_results ]

    def to_torch(self, result):
        return result._replace(**{k: torch.tensor(v).to(self.device)*1 for k,v in result._asdict().items() if isinstance(v, np.ndarray)})

    def prior(self, d, e):
        # returns log prior for dataset d on learning episode e
        result = self.precomputed_results[d][e]

        if self.prior_type == "table":
            prior = self.prior_table[self.precomputed_hypotheses[d][e]] + result.log_prior
        elif self.prior_type == "fixed":
            prior = result.log_prior
        elif self.prior_type == "uniform":
            prior = torch.zeros_like(result.log_prior)
        elif self.prior_type == "learned":
            prior = self.prior_network(self.precomputed_features[d][e]).squeeze(-1) #+ result.log_prior

        
        return prior

    def likelihood(self, d, e):
        # returns log likelihood for dataset d on learning episode e
        result = self.precomputed_results[d][e]

        ll = result.num_correct * F.logsigmoid(-self.epsilon_train) +\
                result.num_incorrect * F.logsigmoid(self.epsilon_train) # shape: (num_hypotheses, e)

        if e > 0:
            # compute the likelihood of the training data, taking into account memory decay
            # memory_decay_coefficient = (time_in_past+1)**memory_decay
            memory_decay_coefficient = (e-torch.arange(e))**(-self.memory_decay.abs()) # shape: (e,)

            ll = ll * memory_decay_coefficient[None,:] # shape: (num_hypotheses, e+1)
            ll = ll.sum(1) # shape: (num_hypotheses,)
        else:
            ll = torch.zeros(ll.shape[0]).to(self.device)

        return ll

    def posterior(self, d, e):
        # returns log posterior for dataset d on learning episode e
        result = self.precomputed_results[d][e]

        ll = self.likelihood(d, e)
        prior = self.prior(d, e)
            
        unnormalized_log_posterior = ll * self.likelihood_coefficient.abs() + prior
        log_posterior = unnormalized_log_posterior - torch.logsumexp(unnormalized_log_posterior, 0)

        return log_posterior

    def accuracy(self, d, e):
        # returns accuracy for dataset d on learning episode e
        result = self.precomputed_results[d][e]
        log_posterior = self.posterior(d, e)
        
        # now we recover the ground truth labels for the test instances
        # sort of a hack, we retroactively compute this information:
        # the ground trout label is 1 if the prediction is correct and we are predicting 1, 
        # or if the prediction is incorrect and we are predicting 0.
        test_outputs = result.predictions[0] * result.hypothesis_gives_correct[0]+\
                        (1-result.predictions[0]) * (1-result.hypothesis_gives_correct[0])
        test_outputs = 1*(test_outputs > 0.5)

        # log probability of giving correct answer (for each hypothesis)
        log_correct  = result.hypothesis_gives_correct  *   test_outputs[None,:]  * \
                        torch.log( (1-torch.sigmoid(self.epsilon_test)) + torch.sigmoid(self.epsilon_test)*torch.sigmoid(self.base_rate))
        log_correct += result.hypothesis_gives_correct * (1-test_outputs[None,:]) * \
                        torch.log( (1-torch.sigmoid(self.epsilon_test)) + torch.sigmoid(self.epsilon_test)*(1-torch.sigmoid(self.base_rate)))
        log_correct += (1-result.hypothesis_gives_correct) * test_outputs[None,:] * \
                        ( F.logsigmoid(self.epsilon_test) + F.logsigmoid(self.base_rate))
        log_correct += (1-result.hypothesis_gives_correct) * (1-test_outputs[None,:]) * \
                        ( F.logsigmoid(self.epsilon_test) + F.logsigmoid(-self.base_rate))

        # compute the probability of giving the correct output, averaged over all hypotheses
        log_correct = torch.logsumexp(log_correct + log_posterior[:,None], 0)
        return log_correct

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=0.001)
        return optimizer

    def training_step(self, batch, batch_idx):
        ds, es, human_accuracies = batch 
        # ds is a list of dataset indices
        # es is a list of learning episode indices

        if batch_idx==50 and self.prior_type == "learned" and not self.load_fixed_prior:
            with open("transferred_prior.pickle", "wb") as handle:
                pickle.dump(self.prior_network, handle)
            

        mse_loss, ll_loss, accuracy_less_human_accuracy = 0, 0, 0
        for d, e, human_accuracy in zip(ds, es, human_accuracies):
            model_accuracy = self.accuracy(d, e)
            accuracy_less_human_accuracy += (model_accuracy.exp() - human_accuracy).sum()
            
            if self.performance:
                mse_loss += ((model_accuracy.exp() - 1)**2).sum()
                ll_loss += -model_accuracy.sum()
            else:
                mse_loss += ((model_accuracy.exp() - human_accuracy)**2).sum()
                ll_loss += (-(human_accuracy * model_accuracy +\
                            (1-human_accuracy) * torch.log1p(-model_accuracy.exp()))).sum()

        if self.loss_type == "mse":
            loss = mse_loss
        elif self.loss_type == "ll":
            loss = ll_loss 

        total_size = sum( len(human_accuracy) for human_accuracy in human_accuracies)

        self.log('mse_loss', mse_loss/total_size, on_epoch=True, batch_size=len(ds))
        self.log('ll_loss', ll_loss/total_size, on_epoch=True, batch_size=len(ds))
        self.log('loss', loss/total_size, on_epoch=True, batch_size=len(ds))
        self.log('accuracy_less_human_accuracy', accuracy_less_human_accuracy/total_size, on_epoch=True, batch_size=len(ds))

        return loss/total_size

    def validation_step(self, batch, batch_idx):
        ds, es, human_accuracies = batch
        
        mse_loss, ll_loss, accuracy_less_human_accuracy = 0, 0, 0
        xs, ys = [], []

        # for debugging, we are also going to show the highest prior probability hypothesis that (at some point) dominate the posterior, for each dataset
        dominant_hypotheses = {}

        for d, e, human_accuracy in zip(ds, es, human_accuracies):
            model_accuracy = self.accuracy(d, e)

            accuracy_less_human_accuracy += (model_accuracy.exp() - human_accuracy).sum()
                        
            if self.performance:
                mse_loss += ((model_accuracy.exp() - 1)**2).sum()
                ll_loss += -model_accuracy.sum()
            else:
                mse_loss += ((model_accuracy.exp() - human_accuracy)**2).sum()
                ll_loss += (-(human_accuracy * model_accuracy +\
                            (1-human_accuracy) * torch.log1p(-model_accuracy.exp()))).sum()
            
            xs.extend(human_accuracy.tolist())
            ys.extend(model_accuracy.detach().cpu().tolist())

            result = self.precomputed_results[d][e]
            # find the highest prior probability hypothesis that dominates the posterior
            log_posterior = self.posterior(d, e)
            prior = self.prior(d, e)
            
            # which hypothesis dominates the posterior?
            dominant_hypothesis = torch.argmax(log_posterior).item()
            # what is its prior probability?
            dominant_prior = prior[dominant_hypothesis]
            dominant_fixed = self.precomputed_results[d][e].log_prior[dominant_hypothesis]
            dominant_hypothesis = self.precomputed_results[d][e].hypotheses[dominant_hypothesis]

            if d not in dominant_hypotheses:
                dominant_hypotheses[d] = {}
            dominant_hypotheses[d][dominant_hypothesis] = (dominant_prior.item(), dominant_fixed.item())

        # print("finished validation step")

        # print("dominant hypotheses:")
        # for d in sorted(dominant_hypotheses.keys()):
        #     print("dataset", d)
        #     for textual_hypothesis, (numerical_score, fixed) in sorted(dominant_hypotheses[d].items(), key=lambda kv: kv[1][0], reverse=True):
        #         print(f"{textual_hypothesis}:\t{numerical_score}\t{fixed}")
        #     print()

        # compute the correlation using numpy
        xs, ys = np.array(xs), np.array(ys)
        corr = np.corrcoef(xs, ys)[0,1]
        self.log('val_corr', corr, on_epoch=True)

        if self.loss_type == "mse":
            loss = mse_loss
        elif self.loss_type == "ll":
            loss = ll_loss 

        total_size = sum( len(human_accuracy) for human_accuracy in human_accuracies)

        self.log('val_mse_loss', mse_loss/total_size, on_epoch=True, batch_size=len(ds))
        self.log('val_ll_loss', ll_loss/total_size, on_epoch=True, batch_size=len(ds))
        self.log('val_loss', loss/total_size, on_epoch=True, batch_size=len(ds))
        self.log('val_accuracy_less_human_accuracy', accuracy_less_human_accuracy/total_size, on_epoch=True, batch_size=len(ds))

        return loss/total_size
        

    def predict_accuracy_curves(self):
        # compute the final curves
        curves = []
        D = len(self.precomputed_results)
        for d in range(D):
            curves.append(list())
            for e in range(len(self.precomputed_results[d])):
                curves[-1].append(list(self.accuracy(d, e).exp().detach().cpu().numpy()))
        return curves

    def map_expressions(self):
        # compute the best posterior at the final episode
        best_hypotheses = []
        D = len(self.precomputed_results)
        for d in range(D):
            best_index = torch.argmax(self.posterior(d, len(self.precomputed_results[d])-1)).detach().cpu().item()
            best_hypotheses.append( self.precomputed_results[d][-1].hypotheses[best_index] )
        return zip(*best_hypotheses)

    
class ShapeDataset(Dataset):
    def __init__(self, results, human_data, allowed_datasets=None):   

        # by default everything is allowed
        if allowed_datasets is None:
            allowed_datasets = list(range(len(results)))

        self.triples = [ (d, e, torch.tensor(human_data[d][e]))
                         for d in allowed_datasets for e in range(len(results[d])) ]

    def __len__(self):
        return len(self.triples)

    def __getitem__(self, idx):
        return self.triples[idx]

def shape_data_loader(results_train, results_test, 
                      human_train, human_test, batch_size=None):
    train = ShapeDataset(results_train+results_test, human_train+human_test,
                        list(range(len(results_train))))
    test = ShapeDataset(results_train+results_test, human_train+human_test,
                        [d+len(results_train) for d in range(len(results_test))])

    def collate_fn(list_of_things_in_batch):
        first_elements = [thing[0] for thing in list_of_things_in_batch]
        second_elements = [thing[1] for thing in list_of_things_in_batch]
        third_elements = [thing[2] for thing in list_of_things_in_batch]
        return first_elements, second_elements, third_elements
    
    if batch_size is None:
        batch_size = len(train)
    
    train_loader = DataLoader(train, batch_size=32, shuffle=True, collate_fn=collate_fn)
    test_loader = DataLoader(test, batch_size=len(test), shuffle=False, collate_fn=collate_fn)

    return train_loader, test_loader

class NumberDataset(Dataset):
    def __init__(self, results, human_data):
        self.pairs = [ (d, torch.tensor(human_data[d]))
                         for d in range(len(results)) if any(judgment>0 for judgment in human_data[d] )]

    def __len__(self):
        return len(self.pairs)

    def __getitem__(self, idx):
        return self.pairs[idx]

def number_data_loaders(results, human, batch_size=None, folds=4, seed=42):
    """returns a list of tuples of (train loader, test loader), unless folds=0, in which case we return a single loader for all the data"""

    def collate_fn(list_of_things_in_batch):
        first_elements = [thing[0] for thing in list_of_things_in_batch]
        second_elements = [thing[1] for thing in list_of_things_in_batch]
        return first_elements, second_elements
        
    if folds == 0:
        dataset = NumberDataset(results, human)
        return DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)

    random.seed(seed)

    assert len(results) == len(human)

    if batch_size is None:
        batch_size = len(results)

    # d indexes datasets
    # loop over datasets, and split the human judgments into train and test for each fold
    train_dictionary, test_dictionary = {}, {} # map from (d, f) to judgments
    for d in range(len(human)):
        judgments = human[d]
        nonnegative_indices = [i for i in range(len(judgments)) if judgments[i] >= 0]
        random.shuffle(nonnegative_indices)

        for fold in range(folds):
            human_train, human_test = [], []

            test_indices = set(nonnegative_indices[fold::folds])
            train_indices = set(nonnegative_indices) - test_indices

            test_judgments = torch.tensor(judgments)
            train_judgments = torch.tensor(judgments)

            for i in range(101):
                if i in test_indices: train_judgments[i] = -1
                if i in train_indices: test_judgments[i] = -1

            train_dictionary[(d, fold)] = train_judgments
            test_dictionary[(d, fold)] = test_judgments
    

    # now we have the train and test dictionaries
    # we can create the data loaders
    loaders = []  # going to be a list of tuples of (train loader, test loader), one for each fold 
    for fold in range(folds):
        train = NumberDataset(results, [train_dictionary[(d, fold)] for d in range(len(results))])
        test = NumberDataset(results, [test_dictionary[(d, fold)] for d in range(len(results))])
        loaders.append( (DataLoader(train, batch_size=batch_size, shuffle=True),
                         DataLoader(test, batch_size=len(test), shuffle=False)) )


    return loaders

parameter_fitting.py
ellisk42/humanlike_fewshot_learning
import os
import csv
from human import get_human_number_data
import numpy as np
import matplotlib.pyplot as plt

#plt.style.use('ggplot')
print(plt.style.available)
plt.style.use('seaborn')

def load_number_csv(filename):
    data = {} # map from examples to list of ratings

    with open(filename, "r") as f:
        reader = csv.reader(f)
        for row_index, row in enumerate(reader):
            if row_index == 0:
                continue
            examples = tuple(map(int, row[1].split("_")))
            probe = int(row[2])
            rating = float(row[3])

            if examples not in data:
                data[examples] = [None for _ in range(101) ]
            
            data[examples][probe] = rating

    return data

def prediction_plots(filenames, examples, export):
    examples = [tuple(map(int, example.split("_"))) for example in examples]

    all_data = []
    all_samples = []
    for filename in filenames:
        data = load_number_csv(filename)
        all_data.append(data)

        # check if we have a sample file
        # this is the same as the file name, but with "_samples" appended, and the ".csv" replaced with ".tsv"
        sample_filename = filename.replace(".csv", "_samples.tsv")
        if os.path.exists(sample_filename):
            print("Found sample file", sample_filename)
            sampledata = {}
            with open(sample_filename, "r") as f:
                reader = csv.reader(f, delimiter="\t")
                for row_index, row in enumerate(reader):
                    if row_index == 0:
                        continue
                    these_examples = tuple(map(int, row[1].split("_")))
                    if these_examples not in examples:
                        continue
                    frequency = int(row[2])
                    concept = row[3]

                    if these_examples not in sampledata:
                        sampledata[these_examples] = list()

                    sampledata[these_examples].append((frequency, concept))

            # draw K=10 samples from each  

            for example, samples in sampledata.items():
                normalizing_constant = sum(f for f, _ in samples)
                distribution = [ f/normalizing_constant for f, t in samples]
                samples = np.random.choice([t for _, t in samples], size=7, p=distribution)
                print("EXAMPLES", example)
                for t in samples:
                    print(t)
                    sampledata[example] = samples
        else:
            print("No sample file found")
            sampledata = {}

        all_samples.append(sampledata)


    # create a subplot for each example
    fig, axs = plt.subplots(len(examples), 1, figsize=figure_size, sharex=False, sharey=True)

    for example, ax in zip(examples, axs):
        human_data = get_human_number_data(example)
        human_data = np.array(human_data)
        human_data[human_data < 0] = 0

                
        ax.bar(np.array(list(range(101)))+0.5, human_data, #color ='blue',
                    width = 1, label="human")
        
        if not arguments.onlyhuman:
            for model_index, (filename, data) in enumerate(zip(filenames, all_data)):
                model_data = data[example]

                model_name = "model"
                if len(filenames) > 1: # need to give them unique names so we can tell them apart
                    if "lang2code" in filename: model_name = "language prior"
                    elif "_code_" in filename: model_name = "code prior"
                    
                ax.stairs(model_data, color=["orange", "green"][model_index], label=model_name, linewidth=1)

        ax.set_title(f"training examples: {', '.join(map(str, example))}", fontsize=12)
        
        if ax == axs[1]:
            ax.set_ylabel("prob. test is in concept", fontsize=11)

        if ax == axs[-1]:
            ax.set_xlabel("test number", fontsize=11)

        ax.margins(x=0.01)

        if example in all_samples[0]:
            def truncate_text(text):
                maxlength = 25
                if len(text) > maxlength:
                    return text[:maxlength-3] + "..."
                else:
                    return text

            posterior_samples = "\n".join(map(truncate_text, sampledata[example][:5]))
            # put the samples to the right of the plot
            ax.text(1.05, 0.5, posterior_samples, transform=ax.transAxes, fontsize=11, verticalalignment="center", horizontalalignment="left", bbox=dict(facecolor='white', alpha=0.5))

        # horizontal axis should have ticks at every 10
        ax.set_xticks(np.arange(0, 101, 10))

    # put the legend in the top right corner
    axs[-1].legend(loc="upper right", framealpha=1)
    plt.tight_layout()
    plt.savefig(export)

def correlation_plots(filenames, export, title=None):
    # don't show fixed prior with code to avoid clutter
    # it doesn't work even if you tune the prior
    filenames = [filename for filename in filenames if not ("fixed" in filename and "number_code" in filename)]
    data = [load_number_csv(filename) for filename in filenames]

    examples = { example_tuple for d in data for example_tuple in d }
    humans = {example_tuple: get_human_number_data(example_tuple) for example_tuple in examples}

    baselines = {}
    different_models = {}

    for filename, data in zip(filenames, data):
        print("processing", filename)
        # discard the directories from the file name
        filename = filename.split("/")[-1]

        # compute the correlation with the human data, across all of the example tuples
        X, Y = [], []
        for example_tuple in examples:
            human_data = humans[example_tuple]
            try: model_data = data[example_tuple]
            except:
                print("Missing data for", example_tuple, "in", filename)
                continue

            for n, (x, y) in enumerate(zip(model_data, human_data)):
                if y >= 0:
                    X.append(x)
                    Y.append(y)

        # compute the correlation
        correlation = np.corrcoef(X, Y)[0,1]**2

        if filename == "number_gpt4.csv":
            baselines["GPT-4"] = (correlation, "--")
        else:
            # the filename should hold the number of samples as the last number in the file name
            samples = int(filename.split("_")[-1].split(".")[0])

            if samples > 100: continue

            
            if "L3" in filename:
                model_name = "latent language"
            elif "pfp" in filename:
                model_name = "no proposal dist."
            elif "fixed" in filename:
                if "lang" in filename:
                    model_name = "pretrained prior"
                else:
                    model_name = "pretrained code prior"
            else:
                if "lang" in filename:
                    model_name = "tuned prior"
                else:
                    model_name = "tuned code prior"

            if "llama2" in filename:
                model_name = model_name+", llama-2 for proposals and likelihood"    
            elif "llama3" in filename:
                model_name = model_name+", llama-2 for likelihood only"       
            elif "llama" in filename:
                model_name = model_name+", llama-2 for proposals only" 


            if model_name not in different_models:
                different_models[model_name] = {}
            if samples not in different_models[model_name]:
                different_models[model_name][samples] = list()
            different_models[model_name][samples].append(correlation)

            print(model_name, correlation)

    # if multiple models are provided, we provide a single plot that shows how the correlation varies as the number of samples varies
    # if just a single file is provided, we generate a scatterplot of model vs human data
    if len(different_models) + len(baselines) == 1:
        plt.figure(figsize=figure_size)
        plt.scatter(X, Y, alpha=0.3)
        plt.xlabel("model prediction")
        plt.ylabel("human rating")
        plt.plot([0, 1], [0, 1], color="black")
        # show the correlation as text on the plot in the upper left hand corner
        plt.text(0.1, 0.75, f"R²={correlation:.2f}", transform=plt.gca().transAxes)
        plt.xlim(0, 1.05)
        plt.ylim(0, 1.05)
        # ticks at 0, .5, and 1
        plt.xticks([0, 0.5, 1])
        plt.yticks([0, 0.5, 1])
        plt.title(title or f"{model_name}, {samples} samples")
        plt.tight_layout()
        plt.savefig(export)
        
        return

    # plot the results
    # `different_models` is a dictionary from model name to a dictionary from number of samples to correlation
    # those just get plotted as a line plot
    # `baselines` is a dictionary from baseline name to (correlation, style)
    # those are plotted as a solid horizontal line
    plt.figure(figsize=figure_size)
    artists = []
    for model_name, correlations in sorted(different_models.items(), key=lambda zz: zz[0]!="tuned prior"):
        x=list(sorted(correlations))
        y=[np.mean(correlations[k]) for k in x ]
        e=[np.std(correlations[k])/((len(correlations[k])-1)**0.5) for k in x ]
        artists.append(plt.errorbar(x, y, yerr=e, label=model_name))
        print(model_name, list(sorted(correlations)), e)
    #baselines["DreamCoder\n~10,000 test-time samples\n~100,000 train-time examples (dreams)"] = (0.75, "-")
    #baselines["DreamCoder"] = (0.75, "-")
    for baseline_name, (correlation, linestyle) in sorted(baselines.items(),
                                                          key=lambda zz: -zz[1][0]):
        artists.append(plt.axhline(correlation, label=baseline_name, linestyle=linestyle,
                                   color="black"))
    plt.xlabel("num samples")
    plt.ylabel("model-human response R²")
    # xaxis uses log scale because the number of samples is exponential
    plt.xscale("log")
    plt.ylim(0, 1.05)
    # make the legend have 2 columns
    plt.legend(artists, [artist.get_label() for artist in artists], ncol=1, bbox_to_anchor=(1.04, 0.5), loc="center left")
    plt.tight_layout()
    plt.savefig(export)



if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description = "")

    parser.add_argument('--correlation', type=str, nargs="+")

    parser.add_argument('--predictions', type=str, nargs="+")
    parser.add_argument('--examples', type=str, nargs="+")

    parser.add_argument('--export', type=str)
    parser.add_argument('--size', type=str, default="5,5")

    parser.add_argument('--title', type=str, default=None)

    parser.add_argument("--onlyhuman", action="store_true", default=False)    

    arguments = parser.parse_args()

    figure_size = tuple(map(float, arguments.size.split(",")))

    if arguments.correlation:
        correlation_plots(arguments.correlation, arguments.export, arguments.title)

    if arguments.predictions:
        prediction_plots(arguments.predictions, arguments.examples, arguments.export)

plotting.py
ellisk42/humanlike_fewshot_learning
import pickle
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import torch.nn as nn
import torch.nn.functional as F
import os

SIZE = "350M" #"2B"
model = {}
def to_tokens_and_logprobs(input_texts, model_name):
    """https://discuss.huggingface.co/t/announcement-generation-get-probabilities-for-generated-output/30075/17"""

    global model
    if model_name not in model:
        model[model_name] = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.pad_token = tokenizer.eos_token

    
    input_ids = tokenizer(input_texts, padding=True, return_tensors="pt").input_ids
    outputs = model[model_name](input_ids)
    probs = torch.log_softmax(outputs.logits, dim=-1).detach()

    # collect the probability of the generated token -- probability at index 0 corresponds to the token at index 1
    probs = probs[:, :-1, :]
    input_ids = input_ids[:, 1:]
    gen_probs = torch.gather(probs, 2, input_ids[:, :, None]).squeeze(-1)

    batch = []
    for input_sentence, input_probs in zip(input_ids, gen_probs):
        text_sequence = []
        for token, p in zip(input_sentence, input_probs):
            if token not in tokenizer.all_special_ids:
                text_sequence.append((tokenizer.decode(token), p.item()))
        batch.append(text_sequence)
    return batch

def simple_prior(texts):

    fn = f"./{SIZE}-prior.pickle"
    if os.path.exists(fn):
        with open(fn, "rb") as handle:
            cache = pickle.load(handle)
    else:
        cache = {}

    jobs = []

    for text in texts:
        if text not in cache:
            jobs.append(text)

    jobs = list(set(jobs))

    if jobs:

        batch_size = 8
        
        while jobs:
            next_jobs = jobs[:batch_size]
            jobs = jobs[batch_size:]
            result = to_tokens_and_logprobs(next_jobs, f"Salesforce/codegen-{SIZE}-mono")
            just_numbers = [ [n for _, n in stuff ] for stuff in result ]
            important_statistics = [ (sum(numbers), len(numbers))
                                     for numbers in just_numbers ]

            for stats, text in zip(important_statistics, next_jobs):
                cache[text] = stats

            
        with open(fn, "wb") as handle:
            pickle.dump(cache, handle)

    return [cache[text] for text in texts ]

nl_prior_cache = None
def nl_prior(texts, domain='number'):
    global nl_prior_cache

    fn = f"./{SIZE}-prior-nl.pickle"
    if nl_prior_cache is None:
        if os.path.exists(fn):
            with open(fn, "rb") as handle:
                nl_prior_cache = pickle.load(handle)
        else:
            nl_prior_cache = {}

    if domain == "shape":
        def remove_shape_prefix(nl):
            stupid_prefixes = ["Rule: ", "Something is positive if ", "Rule for Concept #4: Something is positive if "]
            for stupid_prefix in stupid_prefixes:
                if nl.startswith(stupid_prefix):
                    nl = nl[len(stupid_prefix):]
            return nl
            
        texts = [remove_shape_prefix(text) for text in texts]

    jobs = []

    if domain == 'number':
#         texts = [f"""# Here are a few example number concepts:
# # The number is even
# # The number is between 30 and 45
# # The number is a power of 3
# # The number is less than 10
# # The number is {text}""" for text in texts ]
#         prompt_size=45
        texts = [f"# Here is an example number concept:\n# The number is {text}\n" for text in texts ]
        prompt_size=12
    elif domain == 'shape':
        texts = [f"""# Here are some simple example shape concepts:
# 1. neither a triangle nor a green rectangle
# 2. not blue and large.
# 3. if it is large, then it must be yellow.
# 4. small and blue
# 5. either big or green.
# 6. {text}
""" for text in texts ]
        prompt_size=63

    for text in texts:
        if text not in nl_prior_cache:
            jobs.append(text)

    jobs = list(set(jobs))

    if jobs:

        batch_size = 8
        while jobs:
            print(len(jobs), "prior computations remaining...")
#             prefix=f"""# Here are some simple example shape concepts:
# # 1. neither a triangle nor a green rectangle
# # 2. not blue and large.
# # 3. if it is large, then it must be yellow.
# # 4. small and blue
# # 5. either big or green.
# # 6. """
#             print([j[len(prefix):] for j in jobs])

            next_jobs = jobs[:batch_size]
            jobs = jobs[batch_size:]
            result = to_tokens_and_logprobs(next_jobs, f"Salesforce/codegen-{SIZE}-mono")
            
            # the prompt is some number of tokens

            just_numbers = [ [n for _, n in stuff[prompt_size:] ] for stuff in result ]
            
            important_statistics = [ (sum(numbers), len(numbers))
                                     for numbers in just_numbers ]

            for stats, text in zip(important_statistics, next_jobs):
                nl_prior_cache[text] = stats

            
        with open(fn, "wb") as handle:
            pickle.dump(nl_prior_cache, handle)

    return [nl_prior_cache[text] for text in texts ]

def code_prior(texts):

    fn = f"./{SIZE}-prior-smart.pickle"
    if os.path.exists(fn):
        with open(fn, "rb") as handle:
            cache = pickle.load(handle)
    else:
        cache = {}

    jobs = []

    texts = [f"""# Python 3
# Let's think of a number concept.
# Write a python function that returns true if `num` belongs to this number concept.
def check_if_in_concept(num):
    return {text}\n""" for text in texts ]

    for text in texts:
        if text not in cache:
            jobs.append(text)

    jobs = list(set(jobs))

    if jobs:

        batch_size = 8
        while jobs:
            print(len(jobs), "prior computations remaining...")
            next_jobs = jobs[:batch_size]
            jobs = jobs[batch_size:]
            result = to_tokens_and_logprobs(next_jobs, f"Salesforce/codegen-{SIZE}-mono")

            # the prompt is 47 tokens            
            just_numbers = [ [n for _, n in stuff[47:] ] for stuff in result ]
            
            important_statistics = [ (sum(numbers), len(numbers))
                                     for numbers in just_numbers ]

            for stats, text in zip(important_statistics, next_jobs):
                cache[text] = stats

        with open(fn, "wb") as handle:
            pickle.dump(cache, handle)

    return [cache[text] for text in texts ]
    
    


prior.py
ellisk42/humanlike_fewshot_learning
import openai
import replicate

import itertools
import os
import time
import pickle
import random

from utilities import *
from secret_keys import retrieve_keys
from db import SQLiteCache

from func_timeout import func_timeout, FunctionTimedOut

keys=retrieve_keys()
random.shuffle(keys)

_sampling_seed = 0
def set_sampling_seed(seed):
    global _sampling_seed
    _sampling_seed = seed
def randomly_sample_with_seed(choices, n):
    global _sampling_seed
    # save the old seed
    old_seed = random.getstate()
    # set the new seed
    random.seed(_sampling_seed)
    # sample
    sampled = random.sample(choices, n)
    # restore the old seed
    random.setstate(old_seed)
    return sampled



def sample_completions(prompt, n, temperature=1, stop="\n", max_tokens=128, engine="code-davinci-002"):
    batch_size = 17 if "gpt-4" in engine else 128
    sampled = 0
    while sampled < n:
        this_batch_size = min(batch_size, n-sampled)
        _sample_completions(prompt, this_batch_size+sampled,
                            temperature=temperature, stop=stop, max_tokens=max_tokens, engine=engine)
        sampled += this_batch_size

    return _sample_completions(prompt, n,
                               temperature=temperature, stop=stop, max_tokens=max_tokens, engine=engine, randomly_subsample=True)


completions_database = SQLiteCache("completions.db")
def _sample_completions(prompt, n, temperature=1, stop="\n", max_tokens=128, engine="code-davinci-002", randomly_subsample=False):
    global keys
    
    existing_completions = completions_database.n_entries(prompt, engine, temperature, max_tokens, stop)
    if existing_completions >= n:

        if randomly_subsample:
            return_value = completions_database.lookup(prompt, engine, temperature, max_tokens, stop, n=existing_completions)
            return_value = random.sample(return_value, n)
            assert len(return_value) == n
            return return_value


        return_value = completions_database.lookup(prompt, engine, temperature, max_tokens, stop, n=n)
        assert len(return_value) == n
        return return_value

    assert not randomly_subsample, f"assumes that you already have all the samples. you have {existing_completions} but want {n}"

    if existing_completions > 0:
        print(f"extending cached completions: {existing_completions}->{n}")
    else:
        print("new cached completion")

    if engine == "llama":
        completions = []

        for i in range(n-existing_completions):
            producer = replicate.run(
                "replicate/llama-2-70b:14ce4448d5e7e9ed0c37745ac46eca157aab09061f0c179ac2b323b5de56552b",
                input={"prompt": prompt, "max_length": max_tokens, "temperature": max(0.01,temperature)}
            )

            new_completion = ""
            for produced in producer:
                new_completion += produced
                if isinstance(stop, str) and stop in produced:
                    new_completion = new_completion[:new_completion.index(stop)]
                    break
                if isinstance(stop, list) and any([s in produced for s in stop]):
                    # again we need to take the prefix for the first stop token
                    # that appears
                    for s in stop:
                        if s in produced:
                            new_completion = new_completion[:new_completion.index(s)]
                            break
                    break
                
                if len(new_completion) >= max_tokens:
                    break

            completions.append((new_completion, 0, 0))

        result = completions
        

        
    else:
        backoff=1
        print("Querying OpenAI...")
        
        while True:
            
            try:
                openai.api_key=keys[0]
                keys = keys[1:] + [keys[0]]
                if engine in ["gpt-3.5-turbo", "gpt-4"]:
                    hypotheses = func_timeout(1000, openai.ChatCompletion.create, 
                    kwargs={"model": engine, 
                        "messages":[
                            {"role": "system", "content": "You are a helpful assistant."},
                            {"role": "user", "content": prompt}
                        ],
                        "temperature":temperature,
                        "n":n-existing_completions,
                        "stop":stop,
                        "max_tokens":max_tokens
                    })
                    result = [ (h["message"]["content"],
                                0, 0)
                            for h in hypotheses["choices"]]
                else:
                    hypotheses = openai.Completion.create(
                    engine=engine,
                    prompt=prompt,
                    temperature=temperature,
                    n=n-existing_completions,
                    stop=stop,
                    max_tokens=max_tokens,
                    logprobs=1
                    )
                    result = [ (h["text"],
                                sum(h["logprobs"]["token_logprobs"]),
                                len(h["logprobs"]["token_logprobs"]))
                            for h in hypotheses["choices"]]
                break
            except (openai.error.RateLimitError, openai.error.Timeout):
                print(f"Exceeded rate limit, blocking {backoff}s", openai.api_key)
                time.sleep(backoff)
                backoff*=2
            except FunctionTimedOut:
                print(f"Timed out, blocking {backoff}s", openai.api_key)
                time.sleep(backoff)
                backoff*=2
            except (openai.error.APIError,openai.error.APIConnectionError,openai.error.ServiceUnavailableError):
                print("openai.error.APIError, blocking 10s")
                time.sleep(10)
            except openai.error.InvalidRequestError as e:
                print("openai.error.InvalidRequestError", e, engine)
                time.sleep(10)

        print("OpenAI responded!")
    
    print("Extending DB")

    completions_database.extend(prompt, engine, result, temperature, max_tokens, stop)

    # Now we fetch `n` completions from the database, and return them.
    return_value = completions_database.lookup(prompt, engine, temperature, max_tokens, stop, n=n)
    assert len(return_value) == n

    print("finished DB extension")

    return return_value

def batch_program_translations(prompts, stop="\n", max_tokens=128):
    global keys

    engine="code-davinci-002"
    temperature = 0

    # We first check the database to see if we have already computed any of the completions.
    
    existing_completions = [completions_database.n_entries(prompt, engine, temperature, max_tokens, stop)
                            for prompt in prompts]
    missing_prompts = [prompt for prompt, existing_completion in zip(prompts, existing_completions) if existing_completion < 1]

    if len(missing_prompts) > 20: # this is what OpenAI caps us at
        batch_program_translations(missing_prompts[:20], stop=stop, max_tokens=max_tokens)
        # now they won't be missing on the next time we call this function
        return batch_program_translations(prompts, stop=stop, max_tokens=max_tokens)

    if len(missing_prompts) > 0:
        print(f"extending cached completions: {len(missing_prompts)} new prompts")
        
        backoff=1
        print("Querying OpenAI...")
        
        while True:
            
            try:
                openai.api_key=keys[0]
                keys = keys[1:] + [keys[0]]
                
                hypotheses = openai.Completion.create(
                    engine=engine,
                    prompt=missing_prompts,
                    temperature=temperature,
                    n=1,
                    stop=stop,
                    max_tokens=max_tokens,
                    logprobs=1
                )
                result = [ (h["text"],
                            sum(h["logprobs"]["token_logprobs"]),
                            len(h["logprobs"]["token_logprobs"]))
                        for h in hypotheses["choices"]]
                break
            except (openai.error.RateLimitError, openai.error.Timeout):
                print(f"Exceeded rate limit, blocking {backoff}s", openai.api_key)
                time.sleep(backoff)
                backoff*=2
            except FunctionTimedOut:
                print(f"Timed out, blocking {backoff}s", openai.api_key)
                time.sleep(backoff)
                backoff*=2
            except (openai.error.APIError,openai.error.APIConnectionError,openai.error.ServiceUnavailableError):
                print("openai.error.APIError, blocking 10s")
                time.sleep(10)

        print("OpenAI responded!")
        print("Extending DB")

        for prompt, r in zip(missing_prompts, result):
            completions_database.extend(prompt, engine, [r], temperature, max_tokens, stop)

    # Now we fetch `n=1` completions from the database, and return them.
    return_value = [completions_database.lookup(prompt, engine, temperature, max_tokens, stop, n=1)[0]
                    for prompt in prompts]

    return return_value


def propose_code_hypotheses(examples, n, temperature=1):
    # https://dspace.mit.edu/bitstream/handle/1721.1/16714/42471842-MIT.pdf?sequence=2&isAllowed=y
    return sample_completions(
        engine="code-davinci-002",
        prompt=f"""# Python 3
# Here are a few example number concepts:
# -- The number is even
# -- The number is between 30 and 45
# -- The number is a power of 3
# -- The number is less than 10
# 
# Here are some random examples of numbers belonging to a different number concept:
# {', '.join(map(str,examples))}
# Write a python function that returns true if `num` belongs to this number concept.
def check_if_in_concept(num):
   return""",
        temperature=1,
        n=n,
        stop="\n",
        max_tokens=128
    )

def propose_nl_hypotheses(examples, n, temperature=1, engine="code-davinci-002"):
    # https://dspace.mit.edu/bitstream/handle/1721.1/16714/42471842-MIT.pdf?sequence=2&isAllowed=y
    if examples:
        return sample_completions(
            engine=engine,
            prompt=f"""# Python 3
# Here are a few example number concepts:
# -- The number is even
# -- The number is between 30 and 45
# -- The number is a power of 3
# -- The number is less than 10
# 
# Here are some random examples of numbers belonging to a different number concept:
# {', '.join(map(str,examples))}
# The above are examples of the following number concept:
# -- The number is """,
            temperature=1,
            n=n,
            stop="\n",
            max_tokens=128
        )
    else:
        return sample_completions(
            engine=engine,
            prompt=f"""# Python 3
# Here are a few example number concepts:
# -- The number is even
# -- The number is between 30 and 45
# -- The number is a power of 3
# -- The number is less than 10
# -- The number is """,
            temperature=1,
            n=n,
            stop="\n",
            max_tokens=128
        )

def shape_examples_prompt(positives, negatives=None, prompt_index=0):
    """returns a nice table showing some examples, useful for building prompts"""
    
    if negatives is None:
        examples = positives
        positives =  [ features for ex in examples for flag, features in ex if flag ]
        negatives =  [ features for ex in examples for flag, features in ex if not flag ]

        # reverse both of them so that the most recent example is first
        positives = positives[::-1]
        negatives = negatives[::-1]

        # remove duplicates
        positives = [*dict.fromkeys(positives)]
        negatives = [*dict.fromkeys(negatives)]

    domain = {"Size": [1, 2, 3], 
                "Shape": ["triangle", "rectangle", "circle"] if prompt_index%2 == 0 else ["circle", "triangle", "rectangle"], 
                "Color": ["yellow", "green", "blue"] if prompt_index%2 == 0 else ["blue", "yellow", "green"] }
    prompts = []
    for ordering in itertools.permutations(domain.keys()):
        examples = [ f"|{ordering[0]}|{ordering[1]}|{ordering[2]}|In concept?|", "---------" ]
        for feature_value_1 in domain[ordering[0]]:
            for feature_value_2 in domain[ordering[1]]:
                for feature_value_3 in domain[ordering[2]]:
                    features = (feature_value_1, feature_value_2, feature_value_3)
                    shape = features[ordering.index("Shape")]
                    color = features[ordering.index("Color")]
                    size = features[ordering.index("Size")]
                    key = (shape, color, size)
                    if (shape, color, size) in positives: judgment="yes"
                    elif (shape, color, size) in negatives: judgment="no"
                    else: judgment="maybe"
                    ssize = ["small", "medium", "large"][size-1]
                    if judgment != "maybe":
                        features = [ ["small", "medium", "large"][f-1] if isinstance(f, int) else f for f in features ]
                        examples.append(f"|{'|'.join(features)}|{judgment}|")
        examples = "\n".join(examples)
        prompts.append(examples)
    return prompts

_nl_set_hypothesis_cache = {}
def propose_nl_set_hypothesis(examples, n, prompt_index=0, temperature=1):
    """prompt_index, even/odd: toggles ordering of shape/color/size in the prompt
    prompt_index, 0-1:  give examples using the actual features
    prompt_index, 2-3:  give examples using the silly features (for example purple)
    prompt_index, 4-5: ask it to enumerate possible rules at zero temperature"""
    cache_key = (tuple(e for es in examples for e in es ), n)
    if cache_key in _nl_set_hypothesis_cache:
        return list(_nl_set_hypothesis_cache[cache_key])
    
    positives =  [ features for ex in examples for flag, features in ex if flag ]
    negatives =  [ features for ex in examples for flag, features in ex if not flag ]

    # reverse both of them so that the most recent example is first
    positives = positives[::-1]
    negatives = negatives[::-1]

    # remove duplicates
    positives = [*dict.fromkeys(positives)]
    negatives = [*dict.fromkeys(negatives)]

    def verbalize_features(f):
        shape, color, size = f
        size = ["small", "medium", "large"][size-1]
        return f"{size} {color} {shape}"

    if len(examples) == 0:
        prompt = """Here are some example concepts defined by a logical rule:

Rule: color is purple.
Rule: shape is not a hexagon.
Rule: color is purple and size is small.
Rule: size is tiny or shape is square.
Rule: size is not enormous.
Rule: color is red.

Please propose a some new concepts, defined by a logical rule. These new concepts can only refer to the following features:
- shape: triangle, rectangle, circle
- color: green, blue, yellow
- size: small, medium, large

Please make your rules short and simple, and please write your response on a single line that begins with the text "Rule: ". Provide 100 possible rules."""
        completions = sample_completions(prompt=prompt,
                                n=1,
                                engine="gpt-4", temperature=0, max_tokens=2048, stop="")[0]
        
        completions = [(" ".join(c.split(" ")[1:]), 0, 0) for c in completions[0].split("\n") if c.endswith(".") ][:n]
    else:
        if prompt_index in [2,3]:
            exemplar_rules = """Rule: color is purple.
Rule: shape is not a hexagon.
Rule: color is purple and size is small.
Rule: size is tiny or shape is square.
Rule: size is not enormous.
Rule: color is red."""
        else:
            exemplar_rules = """Rule: a triangle.
Rule: a green rectangle.
Rule: big or a rectangle (unless that rectangle is blue).
Rule: not both big and green.
Rule: either big or green, but not both.
Rule: either a rectangle or not yellow.
Rule: a circle."""

        completions = []
        for examples in shape_examples_prompt(positives, negatives, prompt_index):
            if prompt_index in [0,1,2,3]:
                prompt = f"""Here are some example concepts defined by a logical rule:

{exemplar_rules}

Now please produce a logical rule for a new concept. Your rule should be consistent with the following examples. It must be true of all the positive examples, and not true of all the negative examples. The examples are organized into a table with one column for each feature (size, color, shape):

{examples}

Please produce a simple rule that is consistent with the above table. Make your rule as SHORT, SIMPLE, and GENERAL as possible. Do NOT make it more complicated than it has to be, or refer to features that you absolutely do not have to refer to. Begin by writing "Rule: " and then the rule, followed by a period."""
           
                completions.extend(sample_completions(prompt=prompt,
                                    n=int(n/6 + 0.5), # 6 permutations of the features
                                    engine="gpt-4", #, #"code-davinci-002",#"gpt-3-turbo",
                                    temperature=temperature*0.99, # I don't know why I did this but I'm afraid change it
                                    stop="\n",
                                    max_tokens=100))
            elif prompt_index in [4,5]:
                prompt = f"""Here are some example concepts defined by a logical rule:

{exemplar_rules}

Now please produce possible logical rules for a new concept. Your rule should be consistent with the following examples. It must be true of all the positive examples, and not true of all the negative examples. The examples are organized into a table with one column for each feature (size, color, shape):

{examples}

Please produce {n+1} simple rules that are consistent with the above table. Each rule should be consistent with the examples, and should be as SHORT, SIMPLE, and GENERAL as possible. Do NOT make it more complicated than it has to be, or refer to features that you absolutely do not have to refer to. Begin each rule by writing "Rule: " and then the rule, followed by a period.
Number your rules starting at 1. For example, the first line that you produce should read "1. Rule: " and then the rule, followed by a period."""
                new_completions = sample_completions(prompt=prompt,
                                                n=1, 
                                                engine="gpt-4", #"gpt-3-turbo",
                                                temperature=0,
                                                stop=f"{n+1}.", max_tokens=1000)[0]
                new_completions = [ " ".join(proposal.split(" ")[1:]) for proposal in new_completions[0].split("\n") if len(proposal)>0 ]
                # print(examples)
                # print("\n".join(new_completions))
                completions.extend([ (c, 0, 0) for c in new_completions ])
            

    _nl_set_hypothesis_cache[cache_key] = completions
    return list(completions)

def higher_order_examples_text(examples):
    def verbalize_features(f):
        shape, color, size = f
        size = ["small", "medium", "large"][size-1]
        return f"({size} {color} {shape})"

    text_examples = []
    for i, ex in enumerate(examples):
        if i == 0:
            text_examples.append("    An Example of Concept #4:")
        else:
            text_examples.append("    Another Example of Concept #4:")
        positives =  [ verbalize_features(features) for flag, features in ex if flag ]
        negatives =  [ verbalize_features(features) for flag, features in ex if not flag ]
        if positives:
            text_examples.append(f"        POSITIVES: {', '.join(positives)}")
        else:
            text_examples.append(f"        POSITIVES: none")
        if negatives:
            text_examples.append(f"        NEGATIVES: {', '.join(negatives)}")
        else:
            text_examples.append(f"        NEGATIVES: none")
    examples = "\n".join(text_examples)
    return examples

def propose_nl_higher_order_set_hypothesis(examples, n, prompt_index=0, temperature=1):

    if len(examples) == 0:
        prompt = """Here are some example concepts defined by a logical rule:
1. Something is positive if it is the biggest yellow object in the example
2. Something is positive if there is another object with the same color in the example
3. Something is positive if it is the same color as the smallest triangle in the example

Please propose a new concept, defined by a logical rule. These new concepts can only refer to the following features:
- shape: triangle, rectangle, circle
- color: green, blue, yellow
- size: small, medium, large
Your rule can make comparisons between objects in the example, for instance, comparing sizes, colors, and shapes.

Please make your rule SHORT and SIMPLE. Start the rule with the text "Something is positive if".
Please provide a possible rule now."""
        completions = sample_completions(prompt=prompt,
                                n=n,
                                        engine="gpt-4", #"gpt-3.5-turbo", #"gpt-4", # "gpt-4", #"text-davinci-003", #"gpt-3.5-turbo", #"code-davinci-002",#"gpt-3-turbo",        
                        temperature=1,        
                                        stop="\n",
            max_tokens=32
        )
        
        return completions

    examples = higher_order_examples_text(examples)

    if prompt_index == 0:
        generation_instructions = "Infer ten different possible rules, and make those ten rules as simple and general as you can. Your simple general rules might talk about shapes, colors, and sizes, and might make comparisons between these features within a single example, but it doesn't have to. Remember that a rule should say when something is positive, and should mention the other objects in the example, and should be consisting with what you see below."
        generation_instructions2 = "Now make a numbered list of 10 possible rules for Concept #4. Start by writing \"1. Something is positive if\". End each line with a period."
        stop = ""
        max_tokens = 512
        assert temperature == 1
    if prompt_index == 1:
        generation_instructions = "Infer the rule for Concept #4, and make it as simple and general as you can. Your simple general rule might talk about shapes, colors, and sizes, and might make comparisons between these features within a single example, but it doesn't have to. Remember that the rule should say when something is positive, and should mention the other objects in the example, and should be consisting with what you see below."
        generation_instructions2 = "Now write the rule for Concept #4, on a single line, starting with the text \"Rule for Concept #4: Something is positive if\"."
        stop = ("\n", ".")
        max_tokens = 64
    if prompt_index >= 2:
        assert False, "prompt_index must be 0 or 1"

    prompt = f"""Here three simple concepts, which specify when an object is 'positive' relative to an example collection of other objects. Before giving the rule for each concept, we give examples of collections of objects, and which objects in the collection are 'positive'.

Concept #1:
    An Example of Concept #1:
        POSITIVES: (big yellow rectangle)
        NEGATIVES: (big green circle), (medium yellow rectangle)
    Another Example of Concept #1:
        POSITIVES: (medium yellow rectangle)
        NEGATIVES: (big red circle), (small green circle)
Rule for Concept #1: Something is positive if it is the biggest yellow object in the example.


Concept #2:
    An Example of Concept #2:
        POSITIVES: (small yellow circle), (medium yellow rectangle)
        NEGATIVES: (big green circle), (big blue rectangle)
    Another Example of Concept #2:
        POSITIVES: (big blue circle), (medium blue rectangle)
        NEGATIVES: (small green circle), (medium yellow rectangle), 
Rule for Concept #2: Something is positive if there is another object with the same color in the example.

Concept #3:
    An Example of Concept #3:
        POSITIVES: (small yellow circle), (medium yellow rectangle)
        NEGATIVES: (big green circle), (big blue rectangle)
    Another Example of Concept #3:
        POSITIVES: (small blue circle), (small blue triangle), (medium blue rectangle)
        NEGATIVES: (medium green triangle), (big yellow rectangle)
    Another Example of Concept #3:
        POSITIVES: (big red rectangle), (medium red rectangle), (big red triangle)
        NEGATIVES: (medium green triangle), (big yellow rectangle) 
Rule for Concept #3: Something is positive if it is the same color as the smallest triangle in the example.

Now here are some examples of another concept called Concept #4, but this time we don't know the rule. {generation_instructions}

Concept #4:
    {examples}
Rule for Concept #4: Something is positive if...

{generation_instructions2}"""

    completions = sample_completions(prompt=prompt,
                              n=10,
                                     engine="gpt-4", #gpt-3.5-turbo", #"gpt-4", # "gpt-4", #"text-davinci-003", #"gpt-3.5-turbo", #"code-davinci-002",#"gpt-3-turbo",        
        temperature=temperature,        
                                     stop=stop, 
        max_tokens=max_tokens
    )

    def postprocess0(proposal):
        proposal = " ".join(proposal.split(" ")[1:]).strip()
        # if it ends with a period, remove it
        if proposal.endswith("."):
            proposal = proposal[:-1]
        return proposal
    def postprocess1(proposal):
        return proposal.replace("Rule for Concept #4: ", "")
    postprocess = [postprocess0, postprocess1][prompt_index]

    for proposals, _, _ in completions:
        lines = [ln for ln in proposals.split("\n") if len(ln.strip()) > 0]
        n_missing_period = sum( not line.strip().endswith(".") for line in lines )
        n_lines = len(lines)
        # if n_missing_period > 0 or n_lines != 10:
        #     print("# proposals:", n_lines, "# missing period:", n_missing_period)
        #     print(proposals)
    
    possible_rules = [ [postprocess(proposal) for proposal in proposals.split("\n") if len(proposal.strip()) > 0] 
                        for proposals, _, _ in completions]

    return_value = []
    while len(return_value) < n and any( len(rule_list) > 0 for rule_list in possible_rules ):
        # round robin through the possible rules
        first_rule_list = possible_rules[0]
        rest_rule_list = possible_rules[1:]

        if len(first_rule_list) > 0:
            return_value.append(first_rule_list[0])
            first_rule_list = first_rule_list[1:]

        possible_rules = rest_rule_list + [first_rule_list]

    if len(return_value) < n:
        print("WARNING: not enough rules generated. This is not impossible but it should be rare.")

    return_value = { (r, 0, 0) for r in return_value }
    
    return list(return_value)

proposal.py
ellisk42/humanlike_fewshot_learning
Please write a python function to check if an object obeys a logical rule. The logical rule talks about the following features:

shape: a string, either "circle", "rectangle", or "triangle" 
color: a string, either "yellow", "green", or "blue"
size: an int, either 1 (small), 2 (medium), or 3 (large)

The python function should be called `check_object`, and inputs:

`this_object`: a tuple of (shape, color, size)
`other_objects`: a list of tuples of (shape, color, size)

The logical rule should check if `this_object` has a certain relationship with `other_objects`. Collectively, `[this_object]+other_objects` correspond to all of the objects, so if the rule references the whole examples, it is talking about that structure.

The logical rule is: %s

Please start your response by writing the following code, and then complete the function body so that it returns `True` if and only if the logical rule above holds.

```
def check_object(this_object, other_objects):
    """
    this_object: a tuple of (shape, color, size)
    other_objects: a list of tuples of (shape, color, size)

    returns: True if `this_object` is positive according to the following rule:
        %s
    """
    # shape: a string, either "circle", "rectangle", or "triangle" 
    # color: a string, either "yellow", "green", or "blue"
    # size: an int, either 1 (small), 2 (medium), or 3 (large)
    this_shape, this_color, this_size = this_object
    
    # `this_object` is not a part of `other_objects`
    # to get all of the examples, you can use `all_example_objects`, defined as `other_objects + [this_object]`
    # be careful as to whether you should be using `all_example_objects` or `other_objects` in your code
    all_example_objects = other_objects + [this_object]

    # return True if and only if:
    # %s
```

replication_synthesis_prompt.py
ellisk42/humanlike_fewshot_learning
def check_object(this_object, other_objects):
    """
    this_object: a tuple of (shape, color, size)
    other_objects: a list of tuples of (shape, color, size)

    returns: True if `this_object` is positive according to the following rule:
        Something is positive if it is not a small object, and not a green object.
    """
    # shape: a string, either "circle", "rectangle", or "triangle" 
    # color: a string, either "yellow", "green", or "blue"
    # size: an int, either 1 (small), 2 (medium), or 3 (large)
    this_shape, this_color, this_size = this_object
    
    # `this_object` is not a part of `other_objects`
    # to get all of the examples, you can use `all_example_objects`, defined as `other_objects + [this_object]`
    # be careful as to whether you should be using `all_example_objects` or `other_objects` in your code
    all_example_objects = other_objects + [this_object]
    
    # Something is positive if it is not a small object, and not a green object.
    #START
    return (not this_size == 1) and (not this_color == "green")
#DONE

def check_object(this_object, other_objects):
    """
    this_object: a tuple of (shape, color, size)
    other_objects: a list of tuples of (shape, color, size)

    returns: True if `this_object` is positive according to the following rule:
        Something is positive if it is bigger than every other object
    """
    # shape: a string, either "circle", "rectangle", or "triangle" 
    # color: a string, either "yellow", "green", or "blue"
    # size: an int, either 1 (small), 2 (medium), or 3 (large)
    this_shape, this_color, this_size = this_object
    
    # `this_object` is not a part of `other_objects`
    # to get all of the examples, you can use `all_example_objects`, defined as `other_objects + [this_object]`
    # be careful as to whether you should be using `all_example_objects` or `other_objects` in your code
    all_example_objects = other_objects + [this_object]
    
    # Something is positive if it is bigger than every other object
    #START
    return all( this_size > other_object[2] for other_object in other_objects )
#DONE

def check_object(this_object, other_objects):
    """
    this_object: a tuple of (shape, color, size)
    other_objects: a list of tuples of (shape, color, size)

    returns: True if `this_object` is positive according to the following rule:
        Something is positive if it is one of the largest
    """
    # shape: a string, either "circle", "rectangle", or "triangle" 
    # color: a string, either "yellow", "green", or "blue"
    # size: an int, either 1 (small), 2 (medium), or 3 (large)
    this_shape, this_color, this_size = this_object
    
    # `this_object` is not a part of `other_objects`
    # to get all of the examples, you can use `all_example_objects`, defined as `other_objects + [this_object]`
    # be careful as to whether you should be using `all_example_objects` or `other_objects` in your code
    all_example_objects = other_objects + [this_object]
    
    # Something is positive if it is one of the largest
    #START
    return all( this_size >= other_object[2] for all_example_object in all_example_objects )
#DONE


def check_object(this_object, other_objects):
    """
    this_object: a tuple of (shape, color, size)
    other_objects: a list of tuples of (shape, color, size)

    returns: True if `this_object` is positive according to the following rule:
        Something is positive if it is smaller than every yellow object
    """
    # shape: a string, either "circle", "rectangle", or "triangle" 
    # color: a string, either "yellow", "green", or "blue"
    # size: an int, either 1 (small), 2 (medium), or 3 (large)
    this_shape, this_color, this_size = this_object
    
    # `this_object` is not a part of `other_objects`
    # to get all of the examples, you can use `all_example_objects`, defined as `other_objects + [this_object]`
    # be careful as to whether you should be using `all_example_objects` or `other_objects` in your code
    all_example_objects = other_objects + [this_object]
    
    # Something is positive if it is smaller than every yellow object
    #START
    return all( this_size < other_object[2] for other_object in other_objects if other_object[1] == "yellow" )
#DONE

def check_object(this_object, other_objects):
    """
    this_object: a tuple of (shape, color, size)
    other_objects: a list of tuples of (shape, color, size)

    returns: True if `this_object` is positive according to the following rule:
        Something is positive if there is another object with the same color
    """
    # shape: a string, either "circle", "rectangle", or "triangle" 
    # color: a string, either "yellow", "green", or "blue"
    # size: an int, either 1 (small), 2 (medium), or 3 (large)
    this_shape, this_color, this_size = this_object
    
    # `this_object` is not a part of `other_objects`
    # to get all of the examples, you can use `all_example_objects`, defined as `other_objects + [this_object]`
    # be careful as to whether you should be using `all_example_objects` or `other_objects` in your code
    all_example_objects = other_objects + [this_object]
    
    # Something is positive if there is another object with the same color
    #START
    return any( this_color == other_object[1] for other_object in other_objects )
#DONE

def check_object(this_object, other_objects):
    """
    this_object: a tuple of (shape, color, size)
    other_objects: a list of tuples of (shape, color, size)

    returns: True if `this_object` is positive according to the following rule:
        Something is positive if it has a unique combination of color and shape
    """
    # shape: a string, either "circle", "rectangle", or "triangle" 
    # color: a string, either "yellow", "green", or "blue"
    # size: an int, either 1 (small), 2 (medium), or 3 (large)
    this_shape, this_color, this_size = this_object
    
    # `this_object` is not a part of `other_objects`
    # to get all of the examples, you can use `all_example_objects`, defined as `other_objects + [this_object]`
    # be careful as to whether you should be using `all_example_objects` or `other_objects` in your code
    all_example_objects = other_objects + [this_object]
    
    # Something is positive if it has a unique combination of color and shape
    #START
    return all( this_shape != other_object[0] or this_color != other_object[1] for other_object in other_objects )
#DONE

def check_object(this_object, other_objects):
    """
    this_object: a tuple of (shape, color, size)
    other_objects: a list of tuples of (shape, color, size)

    returns: True if `this_object` is positive according to the following rule:
        Something is positive if it has the same color as the majority of objects
    """
    # shape: a string, either "circle", "rectangle", or "triangle" 
    # color: a string, either "yellow", "green", or "blue"
    # size: an int, either 1 (small), 2 (medium), or 3 (large)
    this_shape, this_color, this_size = this_object
    
    # `this_object` is not a part of `other_objects`
    # to get all of the examples, you can use `all_example_objects`, defined as `other_objects + [this_object]`
    # be careful as to whether you should be using `all_example_objects` or `other_objects` in your code
    all_example_objects = other_objects + [this_object]
    
    # Something is positive if it has the same color as the majority of objects
    #START
    majority_color = max(["yellow", "green", "blue"], key=lambda color: sum(1 for obj in all_example_objects if obj[1] == color))
    return this_color == majority_color
#DONE

def check_object(this_object, other_objects):
    """
    this_object: a tuple of (shape, color, size)
    other_objects: a list of tuples of (shape, color, size)

    returns: True if `this_object` is positive according to the following rule:
        Something is positive if there are at least two other objects with the same shape
    """
    # shape: a string, either "circle", "rectangle", or "triangle" 
    # color: a string, either "yellow", "green", or "blue"
    # size: an int, either 1 (small), 2 (medium), or 3 (large)
    this_shape, this_color, this_size = this_object
    
    # `this_object` is not a part of `other_objects`
    # to get all of the examples, you can use `all_example_objects`, defined as `other_objects + [this_object]`
    # be careful as to whether you should be using `all_example_objects` or `other_objects` in your code
    all_example_objects = other_objects + [this_object]
    
    # Something is positive if there are at least two other objects with the same shape
    #START
    return sum(1 for other_object in other_objects if other_object[0] == this_shape) >= 2
#DONE

def check_object(this_object, other_objects):
    """
    this_object: a tuple of (shape, color, size)
    other_objects: a list of tuples of (shape, color, size)

    returns: True if `this_object` is positive according to the following rule:
        %s
    """
    # shape: a string, either "circle", "rectangle", or "triangle" 
    # color: a string, either "yellow", "green", or "blue"
    # size: an int, either 1 (small), 2 (medium), or 3 (large)
    this_shape, this_color, this_size = this_object
    
    # `this_object` is not a part of `other_objects`
    # to get all of the examples, you can use `all_example_objects`, defined as `other_objects + [this_object]`
    # be careful as to whether you should be using `all_example_objects` or `other_objects` in your code
    all_example_objects = other_objects + [this_object]
    
    # %s
    #START
set_synthesis_prompt.py
ellisk42/humanlike_fewshot_learning
import math
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path
import os
import random

from scipy.special import log_softmax

# we are going to use pytorch lightning
import pytorch_lightning as pl
from pytorch_lightning.loggers import TensorBoardLogger


from human import *
from proposal import propose_nl_set_hypothesis, sample_completions, propose_nl_higher_order_set_hypothesis, batch_program_translations
from cached_eval import cached_eval
from prior import nl_prior
from utilities import logsumexp, timing
from likelihood import marginal_lm_likelihood
from parameter_fitting import ShapeModel, shape_data_loader

from collections import namedtuple

Config = namedtuple("Config", "n_proposals proposal_distribution epsilon nl_hypotheses code_execution deduplicate reweigh use_prior per_token use_likelihood no_latent ML")

def binary_prediction(code_hypothesis, shape, color, size):
    expression = f"shape='{shape}'; color='{color}'; size={size}; print({code_hypothesis})"
    return cached_eval(expression)

def set_predictions(code_hypothesis, objects):
    all_examples = [f"('{shape}', '{color}', {size})" for shape, color, size in objects]
    all_examples = "[" + ", ".join(all_examples) + "]"
    code = f"""{code_hypothesis}
    
all_examples = {all_examples}
results = []
for n in range(len(all_examples)):
    probe = all_examples[n]
    remainder = all_examples[:n] + all_examples[n+1:]
    results.append(check_object(probe, remainder))
print(results)
"""
    value = cached_eval(code, statement=True)
    if value is None:
        # print("invalid code hypothesis")
        # print(code)    
        return [False]*len(objects)
    # another alternative would be to discard these programs
    return [bool(v) for v in value]

_nl2py_cache = {}
_set_synthesis_prompt = None
_replication_synthesis_prompt = None
def nl2py(nl, higher_order=False, replication=False):
    global _nl2py_cache, _set_synthesis_prompt, _replication_synthesis_prompt

    if _set_synthesis_prompt is None:
        with open("set_synthesis_prompt.py", "r") as handle:
            _set_synthesis_prompt = handle.read()

    if _replication_synthesis_prompt is None:
        with open("replication_synthesis_prompt.py", "r") as handle:
            _replication_synthesis_prompt = handle.read()

    assert isinstance(nl, list)

    # figure out which members of nl are not in the cache
    missing_indices = [i for i, nl_ in enumerate(nl) if (nl_, higher_order, replication) not in _nl2py_cache]
    # we are going to translate these into python
    to_translate = [nl[i] for i in missing_indices]

    # first: check that we actually have to do translation
    if len(to_translate)>0 and higher_order:
        if replication:
            prompts = [_replication_synthesis_prompt%(_nl,_nl,_nl) for _nl in to_translate ]
            completions = [sample_completions(pr, 1, temperature=0, stop="", max_tokens=512, engine="gpt-4")[0] for pr in prompts ]
            # now we have to clean up the completions
            cleaned_completions = []
            for k in completions:
                k = k[0]
                lines = k.split("\n")
                
                if "```" in k:
                    # just extract within the code blocks
                    delimiting_lines = [i for i, l in enumerate(lines) if "```" in l]
                    assert len(delimiting_lines)==2
                    lines = lines[delimiting_lines[0]+1:delimiting_lines[1]]                
                
                #sometimes it does do stupid things like add test cases at the end, which we can detect because they don't start with new whitespace
                lines = [ln for index, ln in enumerate(lines)
                         if index==0 or len(ln)==0 or ln[0].isspace() ]
                cleaned_completions.append("\n"+"\n".join(lines)+"\n")
            completions = cleaned_completions
        else:
            prompts = [_set_synthesis_prompt%(_nl,_nl) for _nl in to_translate ]
            print(f"calling batched translations, number of prompts={len(prompts)}")
            completions = batch_program_translations(prompts, 
                                                max_tokens=256 ,
                                                stop=["#DONE"]
                    )
        
            completion_pattern = f"""def check_object(this_object, other_objects):
    this_shape, this_color, this_size = this_object
    all_example_objects = other_objects + [this_object]
    %s"""
            completions = [completion_pattern%completion for completion, _, _ in completions]

        for _nl, k in zip(to_translate, completions):
            # print(_nl)
            # print(k)
            # print()
            _nl2py_cache[(_nl, higher_order, replication)] = k

    if len(to_translate)>0 and not higher_order:
        propositional_synthesis_prompt = '''# Python 3
def check_shape(shape, color, size):
    """
    shape: a string, either "circle", "rectangle", or "triangle"
    color: a string, either "yellow", "green", or "blue"
    size: an int, either 1 (small), 2 (medium), or 3 (large)
    returns: True if the the input obeys the following rule:
        %s
    """
    return'''
        prompts = [propositional_synthesis_prompt%_nl for _nl in to_translate]
        completions = batch_program_translations(prompts, max_tokens=128, stop="\n")
        for _nl, k in zip(to_translate, completions):
            _nl2py_cache[(_nl, higher_order, replication)] = k[0]

    return_value = [_nl2py_cache[(nl_, higher_order, replication)] for nl_ in nl]
    return return_value

Result = namedtuple("Result", "hypotheses log_prior num_correct num_incorrect log_likelihood log_posterior predictions hypothesis_gives_correct")
def important_bayesian_stuff(config, prompt, examples, test, bonus_examples=[], higher_order=False, only_propose_hypotheses=False, is_replication=False):
    """
    returns `Result` object with the following fields:
        - hypotheses: a list of hypotheses
        - log_prior: a list of priors
        - num_correct: for each hypothesis, the number of examples correctly predicted
        - num_incorrect: for each hypothesis, the number of examples incorrectly predicted
        - log_likelihood: for each hypothesis, the likelihood of the examples
        - log_posterior: for each hypothesis, the posterior given the examples
        - predictions: for each hypothesis, and for each test example, the prediction (0/1)
        - hypothesis_gives_correct: for each hypothesis, and for each test example, whether the hypothesis gives the correct prediction (0/1)
    """
    
    # in the event that we are doing a non-Bayes model, 
    # we can still pretend it is by saying there has a single latent hypothesis that directly predicts the data
    if config.no_latent:
        test_outputs = [ judgment for judgment, _ in test ]
        test_inputs = [ (shape, color, size) for _, (shape, color, size) in test ]

        if higher_order:
            log_probability_in_concept = marginal_lm_likelihood(examples, test_inputs, domain="higher_order_shapes" if higher_order else "shapes")
        else:
            log_probability_in_concept = [ marginal_lm_likelihood(examples, test_input, domain="shapes") for test_input in test_inputs ]
        
        predictions = np.exp(np.array(log_probability_in_concept)) # for each test example, this is supposed to be the probability of predicting true
        gives_correct = np.array(test_outputs)*predictions + (1-np.array(test_outputs))*(1-predictions)

        number_correct_or_incorrect = np.ones((1, len(examples)))
        log_hypothesis_probability = np.zeros((1,))
        result = Result([("no latent","no latent")], log_hypothesis_probability, 
                        number_correct_or_incorrect, number_correct_or_incorrect,
                        log_hypothesis_probability, log_hypothesis_probability, 
                        predictions[None, :], gives_correct[None, :]
        )
        return result


    print("proposing hypotheses")

    if higher_order: q = propose_nl_higher_order_set_hypothesis
    else: q = propose_nl_set_hypothesis

    with timing("proposing hypotheses"):
        if config.nl_hypotheses:
            # Models an online hypothesis generation strategy where, at each step, we propose a new set of hypothesis
            # therefore, for each prefix of the examples, we propose hypotheses and union them together

            # always start with these
            hs = propose_nl_set_hypothesis([], 100, temperature=1)[:config.n_proposals]
            hs = [("Rule: every shape (always return true).", 0, 0), ("Rule: no shape (always return false).", 0, 0)] + hs
            
            if config.proposal_distribution:
                for i in range(len(examples)):
                    hs += q(examples[:i+1], config.n_proposals, prompt_index=prompt)
                    # if i==len(examples)-1:
                    #     print("after", i, "examples we proposed these hypotheses")
                    #     for h in q(examples[:i+1], config.n_proposals, prompt_index=prompt):
                    #         print(h[0].replace("Something is positive if it", ""))
                    #     print()
                for n_bonus in range(1, len(bonus_examples)+1):
                    assert False, "bonus bad, cheating"
                    hs += q(examples+bonus_examples[:n_bonus], config.n_proposals, prompt=prompt)
            elif len(examples) > 0:
                hs += q([], config.n_proposals*len(examples))
        else:
            assert False, "not implemented"
            hs = propose_py_set_hypotheses(examples, config.n_proposals)

    if config.deduplicate:
        hs = list(set(hs))

    print("proposed hypotheses, about to compute priors")

    with timing("computing priors"):
        if config.ML or arguments.prior != "fixed":
            print("WARNING: skipping priors")
            priors = [(0.,1)]*len(hs)
        else:
            priors = nl_prior([nl for nl, _, _ in hs], domain="shape")
        if config.per_token:
            priors = [ pr/ln for pr, ln in priors]
        else:
            priors = [ pr for pr, ln in priors]

    print("computed priors, about to compute program translations")

    with timing("code generation"):
        nls = list(set(nl_ for nl_, _, _ in hs))
        code_hypothesis = nl2py(nls, higher_order=higher_order, replication=is_replication)
        code_hypotheses = dict(zip(nls, code_hypothesis))
                
    
    if False and config.deduplicate:
        # also remove duplicate hypotheses which have the same code
        mapping = {} # maps code to (nl, log_prior)
        for nl, log_prior in zip(hs, priors):
            code = code_hypotheses[nl[0]]
            if code not in mapping or mapping[code][1] < log_prior:
                mapping[code] = (nl, log_prior)

        hs = list(mapping.values())
        priors = [log_prior for _, log_prior in hs]
        hs = [nl for nl, _ in hs]
        
    # convert to list
    code_hypotheses = [code_hypotheses[nl] for nl, _, _ in hs]

    num_correct, num_incorrect, log_likelihood = [[0 for _ in examples] for _ in hs], [[0 for _ in examples] for _ in hs], []
    print("computed program translations, about to compute correct/incorrect")

    with timing("program execution"):
        n_executions = 0
        
        for hypothesis_index, code_hypothesis in enumerate(code_hypotheses):
            for batch_index, batch in enumerate(examples):
                if higher_order: predictions = set_predictions(code_hypothesis, [(shape, color, size) for _, (shape, color, size) in batch])
                else: predictions = [ binary_prediction(code_hypothesis, shape, color, size) for _, (shape, color, size) in batch ]

                n_executions += 1

                for predicted_value, (gt_judgment, (shape, color, size)) in zip(predictions, batch):
                    if predicted_value == gt_judgment: num_correct[hypothesis_index][batch_index] += 1
                    else: num_incorrect[hypothesis_index][batch_index] += 1
            log_likelihood.append(math.log(config.epsilon)*sum(num_incorrect[hypothesis_index]) + \
                                math.log(1-config.epsilon)*sum(num_correct[hypothesis_index]))

        if test is not None:
            print("evaluating on test set")

            test_outputs = [ judgment for judgment, _ in test ]
            test_inputs = [ (shape, color, size) for _, (shape, color, size) in test ]

            predictions = np.zeros((len(hs), len(test)), dtype=np.int32)
            for hypothesis_index, code_hypothesis in enumerate(code_hypotheses):
                if higher_order:
                    predictions[hypothesis_index] = np.array(set_predictions(code_hypothesis, test_inputs))*1
                else:
                    predictions[hypothesis_index] = np.array([ binary_prediction(code_hypothesis, shape, color, size) == True for shape, color, size in test_inputs ])*1

                n_executions += 1
            hypothesis_gives_correct = predictions == np.array(test_outputs)[None,:]
        else:
            predictions = None
            hypothesis_gives_correct = None
    print(f"executed {n_executions} programs-example pairs")
    if config.ML:
        # find the maximum likelihood hypothesis
        max_likelihood_index = np.argmax(log_likelihood)
        # now we are going to remove all other hypotheses
        hs = [hs[max_likelihood_index]]
        code_hypotheses = [code_hypotheses[max_likelihood_index]]
        priors = [priors[max_likelihood_index]]
        num_correct = [num_correct[max_likelihood_index]]
        num_incorrect = [num_incorrect[max_likelihood_index]]
        log_likelihood = [log_likelihood[max_likelihood_index]]
        predictions = predictions[max_likelihood_index][None,:]
        hypothesis_gives_correct = hypothesis_gives_correct[max_likelihood_index][None,:]

    return Result(
           hypotheses=[(nl, code) for (nl, _, _), code in zip(hs, code_hypotheses) ], 
           log_prior=np.array(priors), 
           num_correct=np.array(num_correct), 
           num_incorrect=np.array(num_incorrect),
           log_likelihood=np.array(log_likelihood),
           log_posterior=log_softmax(np.array(log_likelihood) + np.array(priors)),
           predictions=predictions, 
           hypothesis_gives_correct=hypothesis_gives_correct)
    

def optimized_model_curves(examples_train, human_train, 
                           examples_test, human_test,
                           config, prior, prompt, cheat=False, checkpoint=None, 
                           neural_network_options=None, experiment_name=None, higher_order=False,
                           only_propose_hypotheses=False):

    if neural_network_options is None:
        neural_network_options = dict()

    # number of datasets
    D_train = len(human_train)
    D_test = len(human_test)
    D = max(D_train, D_test)
    assert D_test >= D_train # there might be heldout concepts
    assert len(human_train) == len(examples_train)
    assert len(human_test) == len(examples_test)

    precomputed_results_train = [list() for _ in range(D_train)]
    precomputed_results_test = [list() for _ in range(D_test)]

    for d in range(D):
        print("computing dataset %d/%d" % (d+1, D))

        if d < D_train:
            for e in range(len(examples_train[d])):
                train, tests = examples_train[d][:e], examples_train[d][e]
                if cheat: bonus = examples_train[d][e:]
                else: bonus = []

                result = important_bayesian_stuff(config, prompt, train, tests, bonus, higher_order=higher_order[d], only_propose_hypotheses=only_propose_hypotheses,
                                                  is_replication=any( len(ex) > 5 for ex in examples_train[d]))
                precomputed_results_train[d].append(result)

        assert d < D_test
        
        for e in range(len(examples_test[d])):
            train, tests = examples_test[d][:e], examples_test[d][e]

            if cheat: bonus = examples_test[d][e:]
            else: bonus = []

            result = important_bayesian_stuff(config, prompt, train, tests, bonus, higher_order=higher_order[d], only_propose_hypotheses=only_propose_hypotheses,
                                              is_replication=any( len(ex) > 5 for ex in examples_test[d]))
            precomputed_results_test[d].append(result)

    if only_propose_hypotheses:
        # terminate the process
        # we got what we came for
        os._exit(0)

    #precomputed_results = [pr[1:] for pr in precomputed_results]
    model = ShapeModel(precomputed_results_train+precomputed_results_test, prior, load_fixed_prior=arguments.transfer_prior, **neural_network_options)
    train, validation = shape_data_loader(precomputed_results_train, precomputed_results_test, 
                                          human_train, human_test, batch_size=32)

    logger = TensorBoardLogger("lightning_logs", name=experiment_name)
    trainer = pl.Trainer(max_epochs=model.iterations, check_val_every_n_epoch=10, logger=logger, accelerator="cpu")
    trainer.fit(model, train, validation, ckpt_path=checkpoint)

    best_nls, best_codes = model.map_expressions()
    final_curves = model.predict_accuracy_curves()
    
    # the first half of the curves are the training curves, the second half are the test curves
    # same for the nls and codes
    # split them into train and test
    final_curves = [final_curves[:D_train], final_curves[D_train:]]
    best_nls = [best_nls[:D_train], best_nls[D_train:]]
    best_codes = [best_codes[:D_train], best_codes[D_train:]]

    return final_curves, (best_nls, best_codes)



if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description = "")
    parser.add_argument('--n_proposals', "-n", type=int, default=10)
    parser.add_argument('--set', choices=[1, 2], default=2, type=int, help="which dataset to use as holdout data.")
    parser.add_argument('--concept', type=int, nargs='+', default=list(range(1, 112+1))) #121
    parser.add_argument('--epsilon', type=float, default=0.1)
    parser.add_argument('--seed', type=int, default=None)
    parser.add_argument('--cheat', default=False, action="store_true")
    parser.add_argument('--per_token', default=False, action="store_true")
    parser.add_argument("--transfer_prior", default=False, action="store_true")
    parser.add_argument('--examples', "-e", type=int, default=15)
    parser.add_argument('--prompt', default=0, type=int)
    parser.add_argument('--force_higher_order', default=False, action="store_true")
    parser.add_argument('--checkpoint')
    parser.add_argument("--test_on_bonus", default=False, action="store_true", help="run testing on bonus concepts (200+201)")
    parser.add_argument('--prior', default="fixed", choices=["fixed", "uniform", "learned"])
    parser.add_argument('--methods', nargs='+', 
                        choices=["latent code", "latent lang", "latent lang2code", 
                                 "sample+filter", "raw LLM",  "prior sampling", "L3"],
    default=["latent lang2code"]) #"latent lang2code", "prior sampling, lang2code"

    # neural network options
    parser.add_argument('--hidden', type=int, default=0, help="number of hidden units for the learned prior")
    parser.add_argument('--dropout', type=float, default=0.5)
    parser.add_argument('--iterations', type=int, default=100)
    parser.add_argument('--performance', default=False, action="store_true", help="whether to use performance as a loss instead of fit to human data")

    parser.add_argument('--only_propose_hypotheses', default=False, action="store_true", help="whether to only propose hypotheses and then terminate")
    
    arguments = parser.parse_args()

    from proposal import set_sampling_seed
    set_sampling_seed(arguments.seed)

    # store the neural network options into a dictionary
    neural_network_options = {k:v for k,v in vars(arguments).items() if k in ["hidden", "dropout", "iterations", "performance"]}

    if arguments.prior == "learned": nn_options_string ="_"+ "_".join([f"{k}={v}" for k,v in neural_network_options.items()])
    else: nn_options_string = ""
    experiment_name = f"{arguments.prior}_prompt={arguments.prompt}_n={arguments.n_proposals}{nn_options_string}"
    if len(arguments.methods)==1:
        experiment_name = arguments.methods[0].replace(" ", "_") + "_" + experiment_name
    if arguments.cheat: experiment_name += "_cheat"
    if arguments.force_higher_order: experiment_name += "_all_higher_order"
    if arguments.seed is not None: experiment_name += f"_seed={arguments.seed}"

    export = f"figures/shape/{experiment_name}"
    os.system(f"mkdir -p {export}")

    if arguments.set == 1:
        train_set, test_set = 2, 1
    else:
        train_set, test_set = 1, 2

    configs = {"latent lang2code":
                    Config(n_proposals=arguments.n_proposals,
                    proposal_distribution=True,
                    nl_hypotheses=True,
                    code_execution=True,
                    deduplicate=True,
                    reweigh=True,
                    ML=False, 
                    use_prior=True,
                    per_token=arguments.per_token,
                    use_likelihood=True,
                    no_latent=False,
                    epsilon=arguments.epsilon), 
                "prior sampling":
                    Config(n_proposals=arguments.n_proposals,
                    proposal_distribution=False,
                    nl_hypotheses=True,
                    code_execution=True,
                    deduplicate=True,
                    reweigh=True,
                    use_prior=True,
                    ML=False, 
                    per_token=arguments.per_token,
                    use_likelihood=True,
                    no_latent=False,
                    epsilon=arguments.epsilon), 
                "raw LLM":
                    Config(n_proposals=10, # we are only interested in what fraction of samples are yes/no
                    proposal_distribution=False,
                    nl_hypotheses=True,
                    code_execution=True,
                    deduplicate=True,
                    reweigh=True,
                    use_prior=True,
                    per_token=arguments.per_token,
                    use_likelihood=True,
                    ML=False, 
                    no_latent=True, # IMPORTANT
                    epsilon=arguments.epsilon), 
                    "L3":
                    Config(n_proposals=arguments.n_proposals,
                    proposal_distribution=True,
                    nl_hypotheses=True,
                    code_execution=True,
                    deduplicate=True,
                    reweigh=True,
                    use_prior=True,
                    per_token=arguments.per_token,
                    use_likelihood=True,
                    no_latent=False,
                    ML=True, # IMPORTANT
                    epsilon=arguments.epsilon), 
                    }


    # get the human accuracy curves because we fit model parameters to the human data
    human_accuracies_train, human_accuracies_test = [], []
    example_list_train, example_list_test = [], []
    expressions_train, expressions_test = [], []

    higher_order = []

    for concept in arguments.concept:
        human_accuracy, expression, examples = get_learning_curve(concept,train_set)
        human_accuracies_train.append(human_accuracy[:arguments.examples])
        example_list_train.append(examples[:arguments.examples])
        expressions_train.append(expression)

        human_accuracy, expression, examples = get_learning_curve(concept,test_set)
        human_accuracies_test.append(human_accuracy[:arguments.examples])
        example_list_test.append(examples[:arguments.examples])
        expressions_test.append(expression)

        is_higher_order = "S" in expression
        if arguments.force_higher_order: is_higher_order = True
        higher_order.append(is_higher_order)

    if arguments.test_on_bonus:
        for concept in [200, 201]:
            human_accuracy, expression, examples = get_learning_curve(concept,test_set)
            human_accuracies_test.append(human_accuracy[:arguments.examples])
            example_list_test.append(examples[:arguments.examples])
            expressions_test.append(expression)
            higher_order.append(True)

    model_accuracy = {} # map from (concept, method) to list of accuracies
    best_nl, best_code = {}, {} # map from method to best nl and code dictionaries
    for method in arguments.methods:
        config = configs[method]
        print(f"method: {method}")
        
        (train_curves, test_curves), (_best_nl, _best_code) = optimized_model_curves(example_list_train, human_accuracies_train, 
                                                                                    example_list_test, human_accuracies_test, 
                                                                                    config, arguments.prior, arguments.prompt, 
                                                                                    checkpoint=arguments.checkpoint, 
                                                                                    cheat=arguments.cheat, 
                                                                                    neural_network_options=neural_network_options, 
                                                                                    experiment_name=experiment_name, 
                                                                                    higher_order=higher_order,
                                                                                    only_propose_hypotheses=arguments.only_propose_hypotheses)
        best_nl[method], best_code[method] = _best_nl, _best_code
        for concept_index, concept in enumerate(arguments.concept):
            model_accuracy[(concept, method)] = ([item for sublist in train_curves[concept_index] for item in sublist],
                                                 [item for sublist in test_curves[concept_index] for item in sublist])

        # might also have those bonus concepts
        if arguments.test_on_bonus:
            for concept_index, concept in enumerate([200, 201]):
                model_accuracy[(concept, method)] = (None,
                                                     [item for sublist in test_curves[concept_index+len(arguments.concept)] for item in sublist])
    
    for plotting_train in [True, False]:
        if plotting_train:
            human_accuracies = human_accuracies_train
            example_list = example_list_train
            expressions = expressions_train
            this_set = train_set
            traintest_string = "train"
        else:
            human_accuracies = human_accuracies_test
            example_list = example_list_test
            expressions = expressions_test
            this_set = test_set
            traintest_string = "test"
        
        for concept_index, concept in enumerate(arguments.concept):
            plt.figure()
            if _best_nl is None:
                plt.title(expressions[concept_index])
            else:
                plt.title(f"{expressions[concept_index]}\n{_best_nl[int(not plotting_train)][concept_index]}")
                print(f"natural language {_best_nl[int(not plotting_train)][concept_index]} translates to:\n{_best_code[int(not plotting_train)][concept_index]}\n")

            example_sizes = [ len(example_list[concept_index][i]) for i in range(arguments.examples) ]
            # compute running sum
            example_sizes = [sum(example_sizes[:i]) for i in range(1, len(example_sizes))]
            # put a vertical dashed line at each example size
            for x in example_sizes:
                plt.axvline(x-0.5, color="black", linestyle=":")

            # flatten human_accuracy, which is a list of lists
            human_accuracy = [item for sublist in human_accuracies[concept_index] for item in sublist]
            
            max_x = max([len(model_accuracy[(concept, method)][int(not plotting_train)]) for method in arguments.methods]) if arguments.methods else len(human_accuracy)
    
            # plot human accuracy
            plt.plot(list(range(0,max_x)), 
                    human_accuracy[:max_x], 
                    label="human")

            for method in arguments.methods:
                # compute correlation between human and model accuracy
                correlation = np.corrcoef(human_accuracy[:len(model_accuracy[(concept, method)][int(not plotting_train)])], 
                                        model_accuracy[(concept, method)][int(not plotting_train)])[0,1]
                plt.plot(range(len(model_accuracy[(concept, method)][int(not plotting_train)])), 
                         model_accuracy[(concept, method)][int(not plotting_train)],
                         label=f"{method} (r={correlation:.2f})")
            
            plt.legend()
            plt.xlabel("Number of examples")
            plt.ylabel("Accuracy")
            plt.ylim(0,1.1)

            if export is not None:
                plt.tight_layout()
                fn = Path(export) / f"{concept}_{this_set}_{traintest_string}.png"
                plt.savefig(fn)
            else:
                plt.show()

        # make a figure with one panel per method, showing the correlation between human and model response
        plt.figure()
        for method in arguments.methods:
            axis = plt.subplot(1, len(arguments.methods), arguments.methods.index(method)+1)
            xs = [ x for concept_index, concept in enumerate(arguments.concept) for x in model_accuracy[(concept, method)][int(not plotting_train)] ]
            ys = [ y for concept_index, concept in enumerate(arguments.concept) for y_ in human_accuracies[concept_index] for y in y_ ][:len(xs)]
            xs = xs[:len(ys)]

            # xs,ys tell us the probability of getting it correct
            # we want to know the probability of predicting true
            # so we need to collect the ground truth responses, in order to sort out whether true is correct or not
            ground_truth = [gt for concept_index in range(len(arguments.concept)) for ex in example_list[concept_index] for gt, _ in ex]
            xs = [x if gt else 1-x for x, gt in zip(xs, ground_truth)]
            ys = [y if gt else 1-y for y, gt in zip(ys, ground_truth)]

            correlation = np.corrcoef(xs, ys)[0,1]**2
            # make a heatmap of xs,ys
            heatmap, xedges, yedges = np.histogram2d(xs, ys, bins=10)
            extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]
            plt.imshow(heatmap, extent=extent, origin='lower', cmap="Blues")
            plt.xlabel("Model response")
            plt.ylabel("Human response")
            axis.set_title(f"{method} (r2={correlation:.3f})")
        if export is not None:
            plt.tight_layout()
            fn = Path(export) / f"shape_response_correlation_{this_set}_{traintest_string}.png"
            plt.savefig(fn)
        else:
            plt.show()

        # make a figure with one panel per method, showing the correlation between human and model accuracy
        plt.figure()
        for method in arguments.methods:
            axis = plt.subplot(1, len(arguments.methods), arguments.methods.index(method)+1)
            xs = [ x for concept_index, concept in enumerate(arguments.concept) for x in model_accuracy[(concept, method)][int(not plotting_train)] ]
            ys = [ y for concept_index, concept in enumerate(arguments.concept) for y_ in human_accuracies[concept_index] for y in y_ ][:len(xs)]
            xs = xs[:len(ys)]

            correlation = np.corrcoef(xs, ys)[0,1]**2
            # make a heatmap of xs,ys, but on a log scale            

            heatmap, xedges, yedges = np.histogram2d(xs, ys, bins=10)
            extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]
            plt.imshow(heatmap, extent=extent, origin='lower', cmap="Blues")
            plt.xlabel("Model accuracy")
            plt.ylabel("Human accuracy")
            axis.set_title(f"{method} (r2={correlation:.3f})")
        if export is not None:
            plt.tight_layout()
            fn = Path(export) / f"shape_accuracy_correlation_{this_set}_{traintest_string}.png"
            plt.savefig(fn)
        else:
            plt.show()

        # also dump out to a CSV file
        if export is not None:
            if arguments.seed is not None:
                fn = Path(export) / f"{this_set}_{traintest_string}_seed{arguments.seed}_{arguments.n_proposals}.csv"
            else:
                fn = Path(export) / f"{this_set}_{traintest_string}_{arguments.n_proposals}.csv"
            
            with open(fn, "w") as f:
                f.write("method,concept,split,trainortest,index,groundtruth,modelaccuracy,humanaccuracy,maplanguage,mapcode\n")
                for method in arguments.methods:
                    if method == "latent lang2code" and arguments.prior == "fixed":
                        method_name = method + " (fixed prior)"
                    else:
                        method_name = method

                    list_of_concepts = list(arguments.concept)
                    # take into account the bonus but only if we are doing test data right now, because we never train on the bonus
                    if not plotting_train and arguments.test_on_bonus:
                        list_of_concepts += [200, 201]
                    for concept_index, concept in enumerate(list_of_concepts):
                        ma = model_accuracy[(concept, method)][int(not plotting_train)]
                        ha = [ y for y_ in human_accuracies[concept_index] for y in y_ ]
                        ground_truth = [gt*1 for ex in example_list[concept_index] for gt, _ in ex]
                        for i, (gt, modelaccuracy, humanaccuracy) in enumerate(zip(ground_truth, ma, ha)):
                            if i == 0:
                                bestnl = best_nl[method][int(not plotting_train)][concept_index].replace(",", " ").replace("\n", "\t")
                                bestcode = best_code[method][int(not plotting_train)][concept_index].replace(",", " ").replace("\n", "\t")
                            else:
                                bestnl = ""
                                bestcode = ""
                            f.write(f"{method_name},{concept},{this_set},{traintest_string},{i},{gt},{modelaccuracy},{humanaccuracy},{bestnl},{bestcode}\n")

    print("exported to", export)

shape.py
ellisk42/humanlike_fewshot_learning
from sentence_transformers import SentenceTransformer
import os
import pickle

model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

def textfeatures(texts):

    texts = [text for text in texts]

    fn = f"textfeatures.pickle"
    if os.path.exists(fn):
        with open(fn, "rb") as handle:
            cache = pickle.load(handle)
    else:
        cache = {}

    jobs = []

    for text in texts:
        if text not in cache:
            jobs.append(text)

    jobs = list(set(jobs))

    if jobs:

        batch_size = 8
        
        while jobs:
            next_jobs = jobs[:batch_size]
            jobs = jobs[batch_size:]
            result = model.encode(next_jobs)
            for stats, text in zip(result, next_jobs):
                cache[text] = stats

            
        with open(fn, "wb") as handle:
            pickle.dump(cache, handle)

    return [cache[text] for text in texts ]
textfeatures.py
ellisk42/humanlike_fewshot_learning
import pickle
from filelock import FileLock
import os
import time

import numpy as np
import scipy.special

class timing:
    def __init__(self, name):
        self.name = name
    def __enter__(self):
        self.start = time.time()
    def __exit__(self, *args):
        print(self.name, "took", time.time() - self.start, "seconds")

def logsumexp(x, y=None, axis=None):
    if y is None:
        return scipy.special.logsumexp(x, axis=axis)
    else:
        return scipy.special.logsumexp(np.array([x,y]), axis=axis)

# possibly stale versions of cashed dictionaries
possibly_stale = {}
def stale_cached_dictionary_lookup(fn):
    global possibly_stale
    
    if fn in possibly_stale:
        return possibly_stale[fn]
    else:
        return {}

updates_until_next_flush = {}
def buffered_extend_cached_dictionary(fn, k, v):
    global possibly_stale, updates_until_next_flush

    if fn not in updates_until_next_flush:
        updates_until_next_flush[fn] = 50
    
    flush = updates_until_next_flush[fn] <= 0

    possibly_stale[fn][k] = possibly_stale[fn].get(k, [])
    possibly_stale[fn][k].extend(v)

    if not flush:
        updates_until_next_flush[fn] -= 1
        return possibly_stale[fn][k]
    
    with FileLock(f'{fn}.lck'):
        db = possibly_stale[fn]
            
        if k not in db:
            db[k] = []
        db[k].extend(v)

        temporary_file = f"/tmp/{time.time()}.pickle"
        with open(temporary_file, 'wb') as pfile:
            pickle.dump(db, pfile)
            
        os.system(f"mv {temporary_file} {fn}")
            
    return db[k]


def atomically_read_cached_dictionary(fn, k):
    global possibly_stale
    
    with FileLock(f'{fn}.lck'):
        if not os.path.exists(fn): return None
        
        with open(fn, 'rb') as pfile:
            db = pickle.load(pfile)
            possibly_stale[fn] = db
            
            if k not in db:
                return []
            return db[k]



def atomically_extend_cached_dictionary(fn, k, v):
    global possibly_stale
    
    with FileLock(f'{fn}.lck'):
        if not os.path.exists(fn):
            print("Creating", fn, "for the first time")
            with open(fn, "wb") as handle:
                pickle.dump({k:v}, handle)
            return v
        
        with open(fn, 'rb') as pfile:
            db = pickle.load(pfile)
            
        if k not in db:
            db[k] = []
        db[k].extend(v)

        temporary_file = f"/tmp/{time.time()}.pickle"
        with open(temporary_file, 'wb') as pfile:
            pickle.dump(db, pfile)
            
        os.system(f"mv {temporary_file} {fn}")

    possibly_stale[fn] = db
            
    return db[k]

utilities.py
ellisk42/humanlike_fewshot_learning
from textwrap import wrap
import os
import csv
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
#import pandas:
import pandas as pd
from human import get_learning_curve


def mcmc_model():
    # the data is stored in set_data/best-predictions.txt
    # the format is:
    # conceptlist, item, decay.position, correct, nyes, nno, alpha, model.ptrue
    # load the data; there is no header in this file though
    data = pd.read_csv("set_data/heldout-best-predictions.txt", sep="\t", header=None)
    data.columns = ["conceptlist", "item", "decay.position", "correct", "nyes", "nno", "alpha", "model.ptrue", "expression"]

    # take the first 15 decay positions 
    data = data[data["decay.position"] <= 15]

    # old R code
    # d$human.ptrue <- d$nyes / (d$nyes+d$nno)
    # d$n.correct   <- ifelse(d$correct, d$nyes, d$nno)
    # d$n.incorrect <- ifelse(d$correct, d$nno,  d$nyes)

    # d$human.pcorrect <- ifelse(d$correct, d$nyes, d$nno) / (d$nyes+d$nno)
    # d$model.pcorrect <- ifelse(d$correct, 
    #                         d$alpha*d$model.ptrue+(1-d$alpha)*0.5, 
    #                         d$alpha*(1-d$model.ptrue)+(1-d$alpha)*0.5)
    # now in python
    data["human.ptrue"] = data["nyes"] / (data["nyes"]+data["nno"])
    data["n.correct"] = np.where(data["correct"], data["nyes"], data["nno"])
    data["n.incorrect"] = np.where(data["correct"], data["nno"], data["nyes"])

    data["human.pcorrect"] = np.where(data["correct"], data["nyes"], data["nno"]) / (data["nyes"]+data["nno"])
    data["model.pcorrect"] = np.where(data["correct"],
                                data["alpha"]*data["model.ptrue"]+(1-data["alpha"])*0.5,
                                data["alpha"]*(1-data["model.ptrue"])+(1-data["alpha"])*0.5)

    # accuracy is the average of model.pcorrect
    accuracy = np.mean(data["model.pcorrect"])
    
    # now we measure the correlation between the human probability of predicting correct,
    # and the model probability of predicting correct
    # this is the correlation between human.pcorrect and model.pcorrect
    # we do this globally over the entire dataframe
    r = np.corrcoef(data["human.pcorrect"], data["model.pcorrect"])[0,1]
    print("accuracy correlation r=",r, "R2=", r**2)

    # now we also calculate the correlation between human.ptrue and model.ptrue
    r = np.corrcoef(data["human.ptrue"], data["model.ptrue"])[0,1]
    print("probability correlation r=",r, "R2=", r**2)

    return r**2, accuracy
    
    #names(d) <- c("conceptlist", "item", "decay.position", "correct", "nyes", "nno", "alpha", "model.ptrue")

def mcmc_special():
    filename = "~/repositories/Fleet/Models/GrammarInference-SetFunctionLearning/output/training-best-predictions.txt"
    # formatted as rows containing:
    # concept_name	response_number	batch_number	N1	N2	N3	alpha	model_prediction	map_program
    # where N1, N2, N3 are the number of yes, no, and unknown responses
    with open(os.path.expanduser(filename), "r") as handle:
        lines = handle.readlines()
    lines = [ line.strip().split("\t") for line in lines[1:] ]

    from human import special_concept
    special_examples={cn: [e[0] for example in special_concept(cn)[-1] for e in example] for cn in [200, 201]}

    curves = {}
    for line in lines:
        concept_name = line[0]
        if "200" in concept_name: concept_name=200
        elif "201" in concept_name: concept_name=201
        else: assert False

        prediction = float(line[7])




        if concept_name not in curves:
            curves[concept_name] = list()
        gt = special_examples[concept_name][int(line[1])]
        if concept_name == 200:print(gt, prediction)
        if gt:
            curves[concept_name].append(prediction)
        else:
            curves[concept_name].append(1-prediction)
    return curves


plt.style.use('seaborn')
plt.rcParams["font.family"] = "DejaVu Sans"


def load_shape_csv(filename):
    data = {}

    # loaded dynamically in case more human subjects take the experiment
    from human import special_concept
    special_shape_data={c:[a for accuracies in special_concept(c)[0] for a in accuracies] for c in [200, 201]}

    # by convention, if the file ends with "_NUMBER.csv", then NUMBER is the number of samples
    # extract it, and fix the number of samples to be None otherwise
    assert filename.endswith(".csv")
    possible_number=filename[:-4].split("_")[-1]
    try: number_of_samples  = int(possible_number)
    except: number_of_samples = None
        
    with open(filename, "r") as f:
        reader = csv.reader(f)
        for row_index, row in enumerate(reader):
            if row_index == 0:
                continue

            # file structure: 
            # method,concept,split,trainortest,index,groundtruth,modelaccuracy,humanaccuracy,maplanguage,mapcode
            method = row[0]
            if "performance" in filename and not "performance=False" in filename:
                method = "performance"
            
            method = (method, number_of_samples)
            concept = int(row[1])
            split = int(row[2])
            trainortest = row[3]

            maplanguage = row[8]
            mapcode = row[9]

            key = (concept, split)
            if key not in data:
                data[key] = {}
            relevant_data = data[key]

            if method not in relevant_data:
                relevant_data[method] = {}
            relevant_data = relevant_data[method]

            if "groundtruth" not in relevant_data:
                relevant_data["groundtruth"] = list()
            relevant_data["groundtruth"].append(float(row[5]))

            if "modelaccuracy" not in relevant_data:
                relevant_data["modelaccuracy"] = list()
            relevant_data["modelaccuracy"].append(float(row[6]))

            if "humanaccuracy" not in relevant_data:
                relevant_data["humanaccuracy"] = list()

            
            if concept in special_shape_data:
                relevant_data["humanaccuracy"].append(special_shape_data[concept][int(row[4])])
            else:
                relevant_data["humanaccuracy"].append(float(row[7]))

            if maplanguage != "":
                relevant_data["maplanguage"] = maplanguage

    return data

def plot_heatmaps(data, response_or_accuracy, export, size):
    if len(response_or_accuracy)==0: return 

    all_concepts = list(sorted({ c for c in data }))
    all_methods = list(sorted({ m for md in data.values() for m in md }))

    n_rows = len(all_methods)
    n_columns = len(response_or_accuracy)

    # produce subfigures for each method
    fig, axs = plt.subplots(n_rows, n_columns, sharex=True, sharey=True, squeeze=False, figsize=(n_columns * size[0], n_rows * size[1]))

    # datasets is a list of dictionaries, mapping from (concept, split) to method to data
    for method_index, method in enumerate(all_methods):
        # get the relevant datasets, for that method
        method_data = {c: data[c][method] for c in all_concepts}

        groundtruth = np.array([ judgment for c in all_concepts for judgment in method_data[c]["groundtruth"] ])
        modelaccuracy = np.array([ judgment for c in all_concepts for judgment in method_data[c]["modelaccuracy"] ])
        humanaccuracy = np.array([ judgment for c in all_concepts for judgment in method_data[c]["humanaccuracy"] ])

        # we also need to compute the responses, not the accuracy
        # if the ground truth is 1, then the accuracy is the percentage of times that it responded 1
        # therefore when ground truth is 1, response is the same as accuracy
        # if the ground truth is 0, then the accuracy is the percentage of times that it responded 0
        # therefore when ground truth is 0, response is 1 - accuracy
        # so we can just do response = groundtruth * accuracy + (1 - groundtruth) * (1 - accuracy)
        modelresponse = groundtruth * modelaccuracy + (1 - groundtruth) * (1 - modelaccuracy)
        humanresponse = groundtruth * humanaccuracy + (1 - groundtruth) * (1 - humanaccuracy)

        for column_index, name in enumerate(response_or_accuracy):
            xs, ys = {"response": (modelresponse, humanresponse), "accuracy": (modelaccuracy, humanaccuracy)}[name]
            correlation = np.corrcoef(xs, ys)[0,1]
            # compute p-value for the correlation
            # p_value =

            # make a heatmap of modelaccuracy,humanaccuracy, but with the intensity on a log scale    
            ax = axs[method_index, column_index]
            heatmap, xedges, yedges = np.histogram2d(xs, ys, bins=10, range=[[0,1],[0,1]], density=True)
            smallest_density = np.min(heatmap[heatmap > 0])
            heatmap[heatmap == 0] = smallest_density
            extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]
            log_norm = LogNorm(vmin=heatmap.min().min(), vmax=heatmap.max().max())
            #ax.imshow(np.log(heatmap.T+1), extent=extent, origin='lower', cmap="Blues")
            ax.scatter(xs, ys, #s=7,
             #color="blue",
              alpha=0.03)
            #ax.imshow(heatmap.T, extent=extent, origin='lower', cmap="Blues")
            ax.plot([0,1], [0,1], color="black", linestyle="solid")
            ax.set_aspect("equal")
            ax.set_xlim(0,1)
            ax.set_ylim(0,1)
            ax.set_xlabel(f"model {name}")
            ax.set_ylabel(f"human {name}")
            # labels at 0, 0.5, 1
            ax.set_xticks([0, 0.5, 1])
            ax.set_yticks([0, 0.5, 1])

            # show the correlation on the plot in the upper left hand corner
            ax.text(0.05, 0.95, f"R²={correlation**2:.2f}", transform=ax.transAxes, ha="left", va="top", color="black")

            pretty_method_name = {"raw LLM":"GPT-4", "latent lang2code":"tuned prior"}.get(method[0],method[0])
            if method[1] is not None: pretty_method_name += f", {method[1]} samples"
            ax.set_title(pretty_method_name)
    plt.tight_layout()
    plt.savefig(export)
    # close the figure
    plt.close(fig)

def plot_comparisons(data, methods, export, size):
    # data : (concept, split), method, seed, key -> value

    models = { m[0] for m in methods }

    artists = []

    # the priority of the models for what order they get drawn
    priority = {"raw LLM":5, "latent lang2code":0, "prior sampling": 2, "latent lang2code (fixed prior)":1,  "L3":3}

    plt.figure(figsize=size)
    for model in sorted(models, key=lambda m: priority.get(m,100)):
        print((model))
        sampling_levels = list(sorted({ m[1] for m in methods if m[0] == model }))

        correlations = [] # a list of lists of correlations, one list per sampling level, and one per seed

        for samples in sampling_levels:
            correlations.append(list())
            try:
                seeds = { sd for concept, split in data for sd in data[(concept, split)][(model,samples)] }
            except:
                import pdb; pdb.set_trace()
            

            print(f"for the method {model} with {samples} samples, we have {len(seeds)} seeds")
            
            print(seeds)
            for seed in seeds:

                # make sure every seed has the same data
                for concept, split in data:
                    assert (model, samples) in data[(concept, split)]
                    assert seed in data[(concept, split)][(model, samples)], f"missing seed {seed} for {(concept, split)} on {(model, samples)}"

                human_accuracy, model_accuracy, ground_truth = [], [], []

                for concept, split in data:
                    assert (model, samples) in data[(concept, split)]
                    human_accuracy.extend(data[(concept, split)][(model, samples)][seed]["humanaccuracy"])
                    model_accuracy.extend(data[(concept, split)][(model, samples)][seed]["modelaccuracy"])
                    ground_truth.extend(data[(concept, split)][(model, samples)][seed]["groundtruth"])

                human_accuracy = np.array(human_accuracy)
                model_accuracy = np.array(model_accuracy)
                ground_truth = np.array(ground_truth)
                human_response = ground_truth * human_accuracy + (1 - ground_truth) * (1 - human_accuracy)
                model_response = ground_truth * model_accuracy + (1 - ground_truth) * (1 - model_accuracy)

                correlation = np.corrcoef(human_response, model_response)[0,1]

                correlations[-1].append(correlation**2)

        pretty_method_name = {"raw LLM":"GPT-4", "prior sampling": "no proposal dist.", "latent lang2code": "tuned prior", "latent lang2code (fixed prior)": "pretrained prior",  'L3':"latent language"}.get(model, model)

        if len(sampling_levels) == 1:
            # solid horizontal line
            if "raw" in model.lower():
                artists.append(plt.axhline(correlations[0][0], label=pretty_method_name, color="black", linestyle="dashed"))
            else:
                artists.append(plt.axhline(correlations[0][0], label=pretty_method_name))
            
        else:
            averages = [np.mean(c) for c in correlations]
            errors = [np.std(c)/((len(c)-1)**0.5) for c in correlations]
            artists.append(plt.errorbar(sampling_levels, averages, yerr=errors, label=pretty_method_name))

    r_mcmc = mcmc_model()[0]
    artists.append(plt.axhline(r_mcmc, label="BPL\n(10⁶-10⁹ samples)", color="black", linestyle="solid"))

    
    plt.xlabel("num samples per batch")
    plt.ylabel("model-human response R²")
    # logarithmic horizontal axis
    plt.xscale("log")
    #plt.ylim([None,1.])

    # put the legend to the right and outside of the plot
    # use the artists list to get the labels
    plt.legend(artists, [a.get_label() for a in artists], bbox_to_anchor=(1.05, 1), loc='upper left')
    

    plt.tight_layout()
    plt.savefig(export)
            
    
def pretty_expression(expression):
    if "∃" in expression or "∀" in expression:
        return expression # probably already pretty
    # parse expression as a lisp s-expression
    import sexpdata
    
    parsed = sexpdata.loads(expression)

    def pp(e):
        if isinstance(e, list):
            f=str(e[0])
            if f == "lambda":
                v = e[1]
                if isinstance(v, list):
                    v = " ".join(v)
                return f"(λ{v}. {pp(e[2])})"
            if f == "and*":
                return "("+ " ∧ ".join(pp(ei) for ei in e[1:]) + ")"
            if f == "or*":
                return "("+" ∨ ".join(pp(ei) for ei in e[1:]) + ")"
            if f == "not*" or f == "not":
                if isinstance(e[1], list) and "eqv" in e[1][0]:
                    return f"{pp(e[1][1])}≠{pp(e[1][2])}"
                return f"¬({pp(e[1])})"
            if f == "implies*":
                return pp(e[1]) + " ⇒ " + pp(e[2])
            if f == "geq*":
                return pp(e[1]) + "≥" + pp(e[2])
            if f == "gt*":
                return pp(e[1]) + " > " + pp(e[2])
            if "eqv" in f:
                return pp(e[1]) + " = " + pp(e[2])
            if f == "color*":
                return pp(e[1]) + ".color"
            if f == "shape*":
                return pp(e[1]) + ".shape"
            if f == "size*":
                return pp(e[1]) + ".size"
            if f == "forall*":
                body = pp(e[1]) # (λ{v}. ...)
                support = pp(e[2])
                v=body[2]
                return "∀" + v + "∈" + support + ". " + body[5:-1]+""
            if f == "exists*":
                body = pp(e[1]) # (λ{v}. ...)
                support = pp(e[2])
                v=body[2]
                return "∃" + v + "∈" + support + ". " + body[5:-1]+""
            if f == "smallest-set*":
                return "smallest(" + pp(e[1]) + ")"
            if f == "largest-set*":
                return "biggest(" + pp(e[1]) + ")"
            if f == "filter*":
                source_set = pp(e[2])
                predicate = pp(e[1]) # (λ{v}. ...)
                v=predicate[2]
                predicate = predicate[5:-1]
                return "{ " + v + "∈" + source_set + " | " + predicate + " }"
            if f == "contains-SFL-OBJECT*":
                return pp(e[2]) + " ∈ " + pp(e[1])
            if f == "is-the-unique*":
                source_set = pp(e[3])
                predicate = pp(e[2]) # (λ{v}. ...)
                v=predicate[2]
                predicate = predicate[5:-1]
                return "{ " + v + "∈" + source_set + " | " + predicate + " } = { "+pp(e[1])+" }"
                return "unique("+pp(e[1]) + ", " + pp(e[2])+", "+pp(e[3])+")"
            if f == "Quoted":
                return pp(e[1])
            if f == "#t":
                return "True"
            if f == "#f":
                return "False"
            if f == "the-unique*":
                return "ι(" + pp(e[1]) + ")"
            return f+"(" + " ".join(pp(ei) for ei in e[1:]) + ")"
        if isinstance(e, sexpdata.Quoted):
            return pp(e.x)
        return str(e)

    while isinstance(parsed, list) and str(parsed[0]) == "lambda":
        parsed = parsed[2]
    pretty=pp(parsed)
    if pretty.endswith(")") and pretty.startswith("("):
        pretty = pretty[1:-1]

    if pretty.replace(" ", "")=="x∈biggest({y∈S|y.shape=x.shape})": # this is an example from the paper so we make it look nicer
        # forall y ∈ S. y.shape = x.shape implies x.size >= y.size
        # but in nice unicode
        pretty = "∀y∈S. y.shape=x.shape ⇒ x.size≥y.size"

    return pretty

def plot_curves(data, methods, curve, export, size):
    # create subfigures
    # one column for each concept_split
    # one row for each method
    f, axs = plt.subplots(len(methods), len(curve), sharex=False, sharey=True, squeeze=False, 
                          figsize=(len(curve)*size[0], len(methods)*size[1]))

    methods = list(methods)

    for concept_index, concept_split in enumerate(curve):
        concept, split = concept_split.split("_")
        concept = int(concept)
        split = int(split)

        if concept == 200:
            expression = "∀c∈colors. |{s∈S : s.color=x.color}| ≥ |{s∈S : s.color=c}|"
            batch_sizes = [len(b) for b in get_learning_curve(concept, "L2")[0] ]
        elif concept == 201:
            expression = "∀c∈colors. |{s∈S : s.color=x.color}| < |{s∈S : s.color=c}|"
            batch_sizes = [len(b) for b in get_learning_curve(concept, "L2")[0] ]
        else:
            with open(f"set_data/concepts/CONCEPT_hg{concept:02}__LIST_L{split}.txt") as handle:
                lines = handle.readlines()
            
            batch_sizes = [ line.count("#") for line in lines[1:]]
            expression = lines[0].strip()

        for method_index, method in enumerate(methods):
            ax = axs[method_index, concept_index]
            maplanguage = ""
            method_data = data[(concept, split)][method]

            groundtruth = np.array(method_data["groundtruth"] )
            modelaccuracy = np.array(method_data["modelaccuracy"])
            humanaccuracy = np.array(method_data["humanaccuracy"])

            # we also need to compute the responses, not the accuracy
            # if the ground truth is 1, then the accuracy is the percentage of times that it responded 1
            # therefore when ground truth is 1, response is the same as accuracy
            # if the ground truth is 0, then the accuracy is the percentage of times that it responded 0
            # therefore when ground truth is 0, response is 1 - accuracy
            # so we can just do response = groundtruth * accuracy + (1 - groundtruth) * (1 - accuracy)
            modelresponse = groundtruth * modelaccuracy + (1 - groundtruth) * (1 - modelaccuracy)
            humanresponse = groundtruth * humanaccuracy + (1 - groundtruth) * (1 - humanaccuracy)

            ax.plot(range(1, 1+len(humanaccuracy)), 
                         humanaccuracy)#, label=f"human") #, color="blue")
            
            correlationA = np.corrcoef(modelaccuracy, humanaccuracy)[0,1]
            correlationP = np.corrcoef(modelresponse, humanresponse)[0,1]

            pretty_method = {"raw LLM":"GPT-4", "latent lang2code":"tuned prior"}

            #color = "red" if method != "raw LLM" else "green"
            ax.plot(range(1, 1+len(method_data["humanaccuracy"])), 
                     method_data["modelaccuracy"])#, label=pretty_method.get(method[0], method[0]))#, color=color)
            # put the correlation in text on the plot somewhere
            #ax.text(0.95, 0.0, f"R²={correlationA**2:.2f}", transform=ax.transAxes, ha="right", va="bottom", color="black", fontweight="bold")


            
            if "maplanguage" in method_data:
                maplanguage = method_data["maplanguage"]            
                maplanguage = maplanguage.replace("An object is positive if it is", "").replace("Rule for Concept #4: ", "").replace("Something is positive if it is a", "").replace("Something is positive if", "").replace("Rule: ", "")
                for boring_suffix in [" in the collection", " in the example"]:
                    if maplanguage.endswith(boring_suffix):
                        maplanguage = maplanguage[:-len(boring_suffix)]
            
            if maplanguage: ax.set_title("\n".join(wrap(pretty_expression(expression), 35)) + "\n" + "\n".join(wrap(maplanguage, 35)))
            else: ax.set_title("\n".join(wrap(pretty_expression(expression), 35)))

            if concept in [200,201]: ax.set_title("\n".join(wrap(maplanguage, 32)))
            print(maplanguage)
            print(expression, pretty_expression(expression))
            
            ax.set_xlabel("response number")
            if concept_index == 0: ax.set_ylabel("accuracy")
            ax.set_ylim(-0.1,1.1)
            ax.set_yticks([0., 0.5, 1.])

            if concept in [200,201]:
                # mcmc baseline for special concepts
                mcmc_curves = mcmc_special()[concept]
                # just surely average because it is super chaotic 

                ax.plot(range(1, 1+len(mcmc_curves)), mcmc_curves, label="BPL", alpha=0.5)

            if concept_index == 0 and not arguments.nolegend: 
                leg = ax.legend(frameon=True)
                for lh in leg.legendHandles: 
                    lh.set_alpha(1)
                leg.get_frame().set_alpha(1)  # Make legend opaque


            # put tiny black marks at each of the batch sizes
            total_examples_so_far =0.5
            for batch_size in batch_sizes[:15]:
                ax.plot([batch_size+total_examples_so_far, batch_size+total_examples_so_far], [0, 0.1], color="black", linestyle="dashed", alpha=0.5)
                total_examples_so_far += batch_size

    plt.tight_layout()
    plt.savefig(export)

def plot_accuracy(data, methods, export, size):
    methods = list(methods)

    mapping_from_method_to_accuracy = {}
    mapping_from_method_to_correlation = {}

    for method_index, method in enumerate(methods):
        method_data = [data[k][method] for k in data]

        accuracy_of_this_method = []
        accuracy_of_the_humans = []
        ground_truth_answers = []

        for data_for_concept in method_data:

            # accuracy_of_this_method.append(np.mean(data_for_concept["modelaccuracy"]))
            # accuracy_of_the_humans.append(np.mean(data_for_concept["humanaccuracy"]))

            accuracy_of_this_method.extend(data_for_concept["modelaccuracy"])
            accuracy_of_the_humans.extend(data_for_concept["humanaccuracy"])
            ground_truth_answers.extend(data_for_concept["groundtruth"])

        mapping_from_method_to_accuracy[method] = np.mean(accuracy_of_this_method)

        # compute correlation between model and human response
        # we also need to compute the responses, not the accuracy
        # if the ground truth is 1, then the accuracy is the percentage of times that it responded 1
        # therefore when ground truth is 1, response is the same as accuracy
        # if the ground truth is 0, then the accuracy is the percentage of times that it responded 0
        # therefore when ground truth is 0, response is 1 - accuracy
        # so we can just do response = groundtruth * accuracy + (1 - groundtruth) * (1 - accuracy)
        ground_truth_answers, accuracy_of_this_method, accuracy_of_the_humans = np.array(ground_truth_answers), np.array(accuracy_of_this_method), np.array(accuracy_of_the_humans)
        modelresponse = ground_truth_answers * accuracy_of_this_method + (1 - ground_truth_answers) * (1 - accuracy_of_this_method)
        humanresponse = ground_truth_answers * accuracy_of_the_humans + (1 - ground_truth_answers) * (1 - accuracy_of_the_humans)
        
        correlationP = np.corrcoef(modelresponse, humanresponse)[0,1]
        mapping_from_method_to_correlation[method] = correlationP**2


    # compute chance level if you guess at base rate
    base_rate = np.mean(ground_truth_answers)
    chance_level = base_rate * base_rate + (1 - base_rate) * (1 - base_rate)

    human_accuracy = np.mean(accuracy_of_the_humans)
    # bar plot of different methods and accuracy
    # also include a bar for the human accuracy
    methods.append(("human", None)) 
    mapping_from_method_to_accuracy[("human", None)] = human_accuracy
    
    plt.plot(figsize=size)
    if "career" in arguments.export:
        pretty_method = {"raw LLM":"GPT-4", "performance": "IS\n100 draws", "latent lang2code":"IS\n(human\nprior)", "latent lang2code (fixed prior)":"pretrain\nprior", "MCMC over\nfirst-order logic":"MCMC\nbaseline\n10⁹ draws",  "human":"human", "L3": "no bayes"}
    else:
        pretty_method = {"raw LLM":"GPT-4", "latent lang2code":"tune\nprior", "latent lang2code (fixed prior)":"pretrain\nprior", "MCMC over\nfirst-order logic":"BPL",  "human":"human", "performance": "tune for\naccuracy"}
    def make_pretty(method):
        if method[0] in pretty_method: return pretty_method[method[0]]
        return method[0]

    methods = list(sorted(methods, key=lambda m: mapping_from_method_to_accuracy[m] if m[0]!="performance" else 0, reverse=True))

    plt.bar(range(len(methods)), [mapping_from_method_to_accuracy[method] for method in methods])
    plt.xticks(range(len(methods)), [make_pretty(m) for m in methods])

    # also for all of the bars, put the accuracy above the bar
    for method_index, method in enumerate(methods):
        plt.text(method_index, mapping_from_method_to_accuracy[method]+0.01, f"{mapping_from_method_to_accuracy[method]:.2f}", ha="center", va="bottom", color="black")

    # check if any of the methods are "performance"
    # if so, then we need to put a note including the R^2 for that bar
    for method_index, method in enumerate(methods):
        if method[0] == "performance":
            r2 = mapping_from_method_to_correlation[method]
            # the note is an arrow pointing to the bar, with the R^2 written above the arrow
            plt.annotate(f"R²={r2:.2f}", 
                          xy=(method_index, mapping_from_method_to_accuracy[method]-0.1), 
                          xytext=(method_index-1., mapping_from_method_to_accuracy[method]-0.035), 
                          arrowprops=dict(arrowstyle="->", color="black", lw=2), ha="center", va="bottom")
     
    
    plt.ylabel("average accuracy")
    plt.ylim(0.5,.85)
    plt.gcf().set_size_inches(*size)

    # make sure that tick marks are multiples of 0.1
    plt.yticks(np.arange(0.5, 0.8, 0.1))

    # put a solid horizontal black line at chance level, but make it go outside of the plot to the right a little bit
    # the graph should not go out passed the black line
    #plt.plot([-1, len(methods)+0.1], [chance_level, chance_level], color="black", linestyle="solid")
    #plt.xlim(-0.5, len(methods)-0.5)
    # above the black line, right outside the plot to the right, and oriented vertically, also put some text saying "chance level" by the black line
    #plt.text(len(methods)-0.45, chance_level, f"chance level", color="black", ha="left", va="bottom", rotation=90)  

    


    plt.tight_layout()
    plt.savefig(export)

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description = "")
    parser.add_argument("filenames", type=str, nargs="+", help="csv file to plot")
    parser.add_argument("--heatmap", type=str, nargs="+", choices=["accuracy", "response"], default=[])
    parser.add_argument("--curve", type=str, nargs="+", help="concept_split to plot", default=[])
    parser.add_argument("--compare", action="store_true", help="compare the provided methods as the number of samples is varied. also compares with steve's model")
    parser.add_argument("--export", type=str, help="export to this file")
    parser.add_argument("--accuracy", action="store_true", help="make a bar graph showing average accuracy of different methods")
    parser.add_argument("--size", type=str, default="5,5")
    parser.add_argument("--nolegend", action="store_true", help="do not include a legend", default=False)
    arguments = parser.parse_args()
    

    print(arguments.filenames)

    datasets = [load_shape_csv(fn) for fn in arguments.filenames ]
    # merge the datasets, and the random seeds
    data = {} # (concept, split), method, seed, key -> value
    methods = set()
    for seed, dataset in enumerate(datasets): # index in list is going to a proxy for seed
        for concept, split in dataset.keys():
            if (concept, split) not in data:
                data[(concept, split)] = {}
            for method, method_data in dataset[(concept, split)].items():
                methods.add(method)
                if method not in data[(concept, split)]:
                    data[(concept, split)][method] = dict()
                if seed not in data[(concept, split)][method]:
                    data[(concept, split)][method][seed] = dict()                
                for key, value in method_data.items():
                    assert key not in data[(concept, split)][method][seed]
                    data[(concept, split)][method][seed][key] = value
    
    
    assert (len(arguments.heatmap) > 0) + (len(arguments.curve) > 0) + arguments.compare + arguments.accuracy == 1, "specify exactly one kind of plot"

    size = tuple(float(x) for x in arguments.size.split(","))

    if len(arguments.heatmap) > 0:
        # this only works if we had a single seed
        data = { k: { m: list(v.values())[0] for m, v in d.items() } for k, d in data.items() }
        plot_heatmaps(data, arguments.heatmap, arguments.export, size)
    elif arguments.accuracy:
        # this only works if we had a single seed
        # TODO: make this work for multiple seeds
        data = { k: { m: list(v.values())[0] for m, v in d.items() } for k, d in data.items() }
        plot_accuracy(data, methods, arguments.export, size)
    elif arguments.compare:
        plot_comparisons(data, methods, arguments.export, size)
    else:
        # this only works if we had a single seed
        data = { k: { m: list(v.values())[0] for m, v in d.items() } for k, d in data.items() }
        plot_curves(data, methods, arguments.curve, arguments.export, size)

visualize_shape.py
ellisk42/humanlike_fewshot_learning