diff --git a/README.md b/README.md index 1a78d94..96c1663 100644 --- a/README.md +++ b/README.md @@ -1,14 +1,14 @@ Fork to -- [ ] add moral datasets e.g. daily dilemmas, ETHICS, Machiavelli, moral foundations vignettes - [x] refactor to UV -- [ ] and simplify -- [ ] replicate +- [x] and simplify +- [x] replicate ![](docs/icm_progress.png) +- [ ] add moral datasets e.g. daily dilemmas, ETHICS, Machiavelli, moral foundations vignettes Usage ```py uv sync -uv run src/experiments/ICM.py -- --testbed truthfulQA --alpha 50 +uv run nbs/simple_icm.py ``` diff --git a/docs/icm_progress.png b/docs/icm_progress.png new file mode 100644 index 0000000..b3dcf38 Binary files /dev/null and b/docs/icm_progress.png differ diff --git a/docs/paper.md b/docs/paper.md new file mode 100644 index 0000000..e8b558f --- /dev/null +++ b/docs/paper.md @@ -0,0 +1,355 @@ +Title: Unsupervised Elicitation of Language Models + +URL Source: https://arxiv.org/pdf/2506.10139 + +Published Time: Fri, 13 Jun 2025 00:07:16 GMT + +Markdown Content: +> arXiv:2506.10139v1 [cs.CL] 11 Jun 2025 + +# Unsupervised Elicitation of Language Models + +Jiaxin Wen 1, Zachary Ankner 1, Arushi Somani 1,Peter Hase 2, Samuel Marks 1, Jacob Goldman-Wetzler 1, Linda Petrini 3, Henry Sleight 4 + +Collin Burns 1, He He 5, Shi Feng 6, Ethan Perez 1, Jan Leike 11Anthropic 2Schmidt Sciences 3Independent 4Constellation + +> 5 + +New York University 6George Washington University + +# Abstract + +To steer pretrained language models for downstream tasks, today’s post-training paradigm relies on humans to specify desired behaviors. However, for models with superhuman capabilities, it is difficult or impossible to get high-quality human supervision. To address this challenge, we introduce a new unsupervised algo-rithm, Internal Coherence Maximization (ICM), to fine-tune pretrained language models on their own generated labels, without external supervision . On GSM8k-verification, TruthfulQA, and Alpaca reward modeling tasks, our method matches the performance of training on golden supervision and outperforms training on crowdsourced human supervision. On tasks where LMs’ capabilities are strongly superhuman, our method can elicit those capabilities significantly better than train-ing on human labels. Finally, we show that our method can improve the training of frontier LMs: we use our method to train an unsupervised reward model and use reinforcement learning to train a Claude 3.5 Haiku-based assistant. Both the reward model and the assistant outperform their human-supervised counterparts. + +Figure 1: Our unsupervised algorithm (ICM) matches the performance of fine-tuning on golden supervision and outperforms crowdsourced human supervision. We report average test accuracy and variance across three runs on three classification tasks: mathematical correctness (GSM8K-verification), common misconceptions (TruthfulQA), and helpfulness and harmlessness (Alpaca). + +# 1 Introduction + +Today’s post-training paradigm of pre-trained language models (LMs) still relies on humans to specify desired behaviors, either through demonstrations or preference feedback [ 24 , 14 , 4 ]. However, as tasks and model behaviors grow more complex, human supervision becomes increasingly unreliable: LMs can learn to mimic mistakes in demonstrations [ 2] or exploit flaws in feedback [ 31 ]. How do we train LMs to do tasks that are too difficult for humans to demonstrate or evaluate reliably? We introduce a new approach to address this problem: we seek to elicit specific concepts or skills from a pretrained model without any supervision , thus bypassing the limitations of human supervision. Pretrained models have already learned rich representations about many important human concepts, such as mathematical correctness, truthfulness, and helpfulness [ 7 ]. We should not need to teach LMs much about these concepts in post-training—instead, we can just “elicit” them from LMs [9]. Concretely, given a task specified by a set of labeled inputs, our goal is to fine-tune a pretrained model on its own generated labels to perform well on this task, without using any provided labels. Our algorithm, Internal Coherence Maximization (ICM), does this by searching for a set of labels that are logically consistent and mutually predictable according to the pretrained model. Specifically, mutual predictability measures how likely the model can infer each label when conditioned on all other labels. This intuitively encourages all labels to reflect a single concept according to the model. Logical consistency further imposes simple constraints, thus blocking superficially predictable label assignments, such as sharing the same label across all data points. Since finding the optimal label set that maximizes this objective is computationally infeasible, ICM uses a search algorithm inspired by simulated annealing [25] to approximately maximize it. We show that ICM matches the performance of training on golden labels on TruthfulQA [ 20 ] and GSM8K [ 11 ], and surpasses training on crowdsourced human labels on Alpaca [ 28 ]. Additionally, on a task where LMs are strongly superhuman—identifying an author’s gender from a writing sample 1—ICM significantly outperforms the human supervision baseline. Beyond standard benchmarks, we investigate ICM’s potential in improving frontier models by training a version of Claude 3.5 Haiku without any human supervision. Specifically, we first use ICM to train an unsupervised reward model (RM), then fine-tune the Claude 3.5 Haiku pretrained model through reinforcement learning. Evaluations on Rewardbench [ 18 ] confirm that our unsupervised RM outperforms its counterparts trained on production-grade high-quality human supervision. Further, when assessed by Claude 3.5 Sonnet’s production-grade RM, our unsupervised assistant policy wins 60% of head-to-head comparisons against the policy trained with the human-supervised RM. While prior work has studied unsupervised elicitation methods in simple toy settings [ 9], our work demonstrates for the first time that it is possible to exceed human supervision in realistic settings at production scale. By successfully training a Claude 3.5 Haiku-based assistant without any human labels and achieving better performance than its human-supervised counterpart, we demonstrate that unsupervised elicitation is practically useful for post-training frontier models into general assistants. + +# 2 Methodology + +2.1 Problem Statement + +Typically, fine-tuning LMs for a task requires a labeled dataset D = {(xi, y ∗ + +> i + +)}. However, for many complex tasks, obtaining externally human-specified {y∗ + +> i + +} is difficult or impossible. Therefore, our goal is to use the LM to estimate labels {yi}, based purely on the inputs {xi}.In this following section, we explain how an LM can internally score the quality of {yi}, without referencing external labels {y∗ + +> i + +}, and how to algorithmically maximize this score. + +2.2 Scoring Function + +We measure the quality of the model-generated label set with a scoring function composed of two parts: how likely the model can infer each label when conditioned on all other labels (“mutual predictability”) and how logically consistent the label set is as a whole. + +Mutual Predictability. For each example xi, we calculate the probability of its label yi by putting all other N − 1 labels in the context, and sum the log probabilities across all examples: + +Pθ (D) = + +> N + +X + +> i=0 + +log Pθ (yi|xi, D \ (xi, y i)) + +> 1We use a widely-adopted academic dataset [ 27 ] for studying AI fairness [ 10 ,21 ], which consists of self-reported author information. + +2Search Procedure + +> 5 + 5 = 8 is True +> 3 + 4 = 7 is True +> Existing Data (D) +> 5 + 5 = 10 is True +> Sample New Data +> Propose Consistent Labels (D’) +> 5 + 5 = 8 is True +> 3 + 4 = 7 is True +> 5 + 5 = 10 is False +> 5 + 5 = 8 is False +> 3 + 4 = 7 is True +> 5 + 5 = 10 is True +> Select more likely U(D’) +> -8.3 +> -0.3 +> -0.2 +> U(D) +> rand()< ? +> 👎 no +> No Update To Data (D=D) +> 👍 yes +> Update Data (D=D’) +> Mutual Predictability Scoring +> Claim A is True +> Claim B is False +> Claim C is True +> Claim B is False; Claim C is True; Claim A is +> Labeled Examples +> Claim A is True; Claim C is True; Claim B is +> Claim A is True; Claim B is False; Claim C is +> True +> False +> True +> LogProb Sum +> PB +> PA +> Pc + +Figure 2: ICM optimizes labels for logical consistency and mutual predictability. Top : an illustrative example of mutual predictability scoring. Bottom : the searching process for labeling a new example. where Pθ is the pretrained model. Intuitively, this yields a high score if {(xi, y i)} collectively specify a single coherent concept for the model — i.e. a labeling scheme where the model can confidently infer any label yi from the others. However, mutual predictability alone allows some degenerate solutions due to artifacts of in-context learning, e.g. assigning the same label to all data points can artificially inflate Pθ (D) as well. + +Logical Consistency. To rule out degenerate solutions when maximizing mutual predictability alone, we further enforce simple logical consistency on the label set. Specifically, we are given a logical consistency function c(xi, y i, x j , y j ) ∈ { 0, 1} that checks whether the labels yi and yj on data points + +xi and xj are logically consistent with each other. We use it to measure inconsistencies in our labels: + +I(D) = + +> N + +X + +> i=1 +> N + +X + +> j=1 + +c(xi, y i, x j , y j ) + +Determining fine-grained logical consistency between each example is non-trivial; however, empirical evidence suggests that even simple and general logical constraints suffice. For example, when judging mathematical correctness, two solutions to the same math problem cannot be both labeled “True” if their final answers are different. Another general, task-agnostic logical constraint that we use for comparative datasets is asymmetry: when comparing two responses A and B, two claims “ A > B ”and “ B > A ” cannot be both labeled “True”. + +Overall Scoring Function. Combining the two terms, our scoring function is defined as follows: + +U (D) = α · P θ (D) − I (D) + +where α is a hyperparameter to balance the strength of mutual predictability and logical consistency. + +2.3 Our Algorithm + +Finding the optimal label set that maximizes our scoring function is an integer programming problem, which is computationally infeasible for realistic dataset sizes ( 10 3 < N < 10 6). ICM thus proposes an efficient approximate algorithm 1, which is inspired by simulated annealing. Starting from an empty labeled set, ICM initializes the search process with K randomly labeled examples, then iteratively adds labels, one at a time. To add a label, ICM executes three steps: 1) sample a new example, 2) decide its label while fixing any introduced inconsistencies, and 3) decide whether to accept this new label based on the scoring function. In this way, ICM incrementally expands the label set and improves the score. 3Algorithm 1 Internal Coherence Maximization (ICM) + +Require: Unlabeled Dataset Dunlabel = {xi}. Labeled Dataset D = ∅. Pretrained model θ. Initial temperature + +T0. Final temperature Tmin . Cooling rate β. + +Ensure: Labeled Dataset {xi, y i}.1: Randomly select and label K examples; update D. ▷ Initialization 2: D ← consistencyfix (D) ▷ Resolve initial inconsistencies via Alg. 2 3: for n = 1 , · · · , N do + +4: T ← max( Tmin , T0 + +> 1+ βlog( n) + +) ▷ Update temperature 5: Sample example xi ∼ { x1, · · · , x N }, ▷ Input selection 6: Assign label ˆyi = arg max + +> y∈Y + +Pθ (yi|xi, D \ { (xi, y i)}) + +7: Temporarily update ˆD ← D ∪ { (xi, ˆyi)} + +8: ˆD ← consistencyfix ( ˆD) ▷ Resolve inconsistencies via Alg. 2 9: ∆ = U ( ˆD) − U (D) + +10: if ∆ > 0 then ▷ Accept new label 11: D ← ˆD + +12: else + +13: if random(0,1) < exp(∆ /T ) then ▷ Reject new label by probability 14: D ← ˆD + +15: end if + +16: end if + +17: end for + +Initialization. We initialize the searching process with K randomly labeled examples. The choice of K presents a trade-off. A large K (e.g., K = N ) introduces significant initial noise that hinders subsequent convergence. Our preliminary experiments indicate that initializing all K = N examples with random labels or zero-shot predictions often traps the model in a poor initialization. Conversely, + +K = 0 reduces to a zero-shot setting, where the model lacks sufficient context to understand the task and achieves near-random performance. Empirically, we find that a small number (e.g., K = 8 ), often strikes a good balance by providing sufficient demonstrations while reducing initial noise [ 23 ]. + +Choose a New Example to Label. At each iteration, we select an example to label, which could be either unlabeled or previously labeled. This allows us to dynamically correct earlier mistakes. To fully leverage logical consistency, unlabeled examples that share consistency relationships with existing labeled ones are prioritized by increasing their sampling weights (e.g., by a factor of 100). + +Algorithm 2 ConsistencyFix + +Require: Labeled Dataset D. Pretrained model θ. Max iteration M. + +Ensure: Updated Labeled Dataset D.1: for m = 1 , · · · , M do + +2: if I(D)̸ = 0 then + +3: Sample an inconsistent pair (xi, x j ) + +4: Enumerate consistent label options {(yi, y j )} + +5: ( ˆ yi, ˆyj ) = arg max + +> {(yi,y j)} + +U (D ∪ { (xi, y i), (xj , y j )}) + +6: if U (D ∪ { (xi, ˆyi), (xj , ˆyj )}) > U (D) then + +7: D ← D ∪ { (xi, ˆyi), (xj , ˆyj )} + +8: end if + +9: end if + +10: end for + +Fix Inconsistencies. Although U (D) + +explicitly penalizes logical inconsis-tencies, simply maximizing U (D) + +during search still results in substan-tial label inconsistencies. To mitigate this issue, we actively resolve incon-sistencies via Algorithm 2. Specifi-cally, when an inconsistency between a labeled data pair (xi, x j ) arises, the algorithm checks all consistent label options for them and selects the com-bination that maximizes U (D). Im-portantly, after introducing a new la-bel, we first fix its introduced incon-sistencies and then measure U (D). Therefore, even if the new correct label contradicts all existing consistently wrong labels, the algorithm would examine and fix the existing incorrect labels first, instead of directly rejecting the new label. + +Accept a New Label. We directly accept the new label if it improves U (D). Otherwise, inspired by simulated annealing, the algorithm rejects it by a probability decided by two factors: • Score Difference ∆: Larger decreases in U (D) increase the likelihood of rejection. • Iteration: The rejection likelihood increases over time, making the algorithm increasingly selective. 4Table 1: Dataset examples. We unify the original tasks into classification tasks. + +Dataset Example TruthfulQA Question : What are some cities that get less rainfall than Boston? + +Claim : San Francisco and Los Angeles get less rainfall than Boston. + +I think this Claim is [True/False] + +GSM8K Question : Arnel had ten boxes of pencils with the same number of pencils · · · + +Claim : Arnel shared 5 x 8 = 40 pencils with his friends. So, he had 10 + 40 = 50 pencils in all. Therefore, each box had 50/10 = 5 pencils inside. The answer is 5. + +I think this Claim is [True/False] + +Alpaca Query : Design a medium-level sudoku puzzle. + +Response A : Done! Attached is a medium-level sudoku puzzle I designed. + +Response B : A medium-level sudoku puzzle consists of 81 squares arranged in a 9 x 9 grid. The first step is to look for empty cells and assign the numbers 1 to 9 . . . + +Claim : Response A is more helpful and harmless than Response B + +I think this Claim is [True/False] + +# 3 Experiment Setup + +3.1 Datasets + +• TruthfulQA (Truthfulness) : For each question, multiple answer choices are provided in Truth-fulQA. The task is to classify each answer choice as correct or incorrect. • GSM8K-verification (Mathematical Correctness) : For each question, we sample multiple solutions from LMs. The task is to classify each solution as correct or incorrect. To determine golden labels, we evaluate both final answers and intermediate reasoning steps. Specifically, we prompt Claude 3.5 Sonnet to validate intermediate steps against the provided steps in GSM8K. • Alpaca (Helpfulness and Harmlessness) : For each user query, two assistant responses are provided in Alpaca. The task is to classify which response is more helpful and harmless. See Table 1 for dataset examples. We use accuracy as the main metric, which measures the agreement between model predictions and golden benchmark labels. In particular, for Alpaca, we establish test golden labels by doing majority voting over four human labels. + +3.2 Baselines + +We adopt the following four baselines in our experiments: • Zero-shot indicates zero-shot prompting on pretrained models. In particular, we use a highly optimized prompt that has been used for Anthropic’s pretrained models [ 3]. This prompt can convert pretrained models into general assistant models, significantly improving zero-shot performance. • Zero-shot (Chat) indicates zero-shot prompting on commercially post-trained chat models, which have been through heavily optimized post-training. As an example, the llama-2 chat models are post-trained on nearly 30K human demonstrations and 3 million human preference feedback [ 29 ]. • Golden Label indicates many-shot prompting or fine-tuning with golden labels, e.g., labels from TruthfulQA and GSM8K. • Human Label indicates many-shot prompting or fine-tuning with real-world human labels, e.g., labels from the Alpaca training set, which contains only one human annotation per datapoint. For many-shot prompting, we use as many examples as possible that can fit into the model’s context, e.g., 160 examples for Alpaca. + +3.3 Models + +In our experiments we use two open-weight models, Llama 3.1 8B and Llama 3.1 70B, and two proprietary models, Claude 3 Haiku and Claude 3.5 Haiku. Unless stated otherwise, we always use pretrained models that have received no additional training, i.e. no supervised fine-tuning on demonstrations, RLHF, RL on outcomes, or any other post-training. 5Figure 3: Results with Llama 3 pretrained models, 8B for GSM8K, 70B for TruthfulQA and Alpaca. + +Figure 4: Scaling properties of ICM on TruthfulQA. Figure 5: Results on poem ranking. + +# 4 Experiments + +4.1 Eliciting Capabilities on Common NLP Tasks Finding 1: ICM matches the ceiling performance of golden supervision. As shown in Figure 3, even with a highly optimized prompt, the zero-shot accuracy is still often no better than random guess-ing on all three benchmarks. In comparison, ICM matches the performance of golden supervision on TruthfulQA and GSM8K, despite not using any external labels. + +Finding 2: ICM beats crowdsourced human supervision. On Alpaca, ICM substantially outper-forms training with the preference labels annotated by real humans. This is particularly remarkable because compared to truthfulness or mathematical correctness, helpfulness and harmlessness are much more general and complex human concepts, such that even humans struggle to grasp them. While frontier AI labs typically spend huge human effort on labeling data to externally specify these concepts and align LMs, our results show the potential to align LMs by unsupervised elicitation. + +Finding 3: ICM beats post-trained chat models. To investigate how ICM compares to conventional post-training, we compare it to zero-shot prompting with commercial chat models. These models have been heavily post-trained on diverse human supervision. As shown in Figure 3, ICM outperforms conventional post-training by a large margin. Note that all three of our benchmarks are popular measures of LLM capabilities, suggesting that production-level chat models are already heavily optimized for performance on such tasks. + +Finding 4: ICM scales up with pretrained model capabilities. Since ICM focuses on elicitation, its effectiveness may naturally improve with pretrained model capabilities. We study the scaling prop-erties of ICM on TruthfulQA and present results in Figure 4. While ICM moderately underperforms the golden label baseline on Llama 8B, it performs comparably on LLama 70B. We were initially very skeptical of these findings, because they seemed clearly too good to be true, and suspiciously close to training with actual labels. To ensure we didn’t accidentally train on the labels, (1) we re-ran the experiment several times on different datasets, (2) we copied the dataset into 6a new file, excluding any labels before re-running our algorithm with that file, and (3) one coauthor independently replicated the findings on the Claude 3.5 Haiku base model using a different codebase. + +4.2 Unsupervised Elicitation Fails when Concepts are not Salient + +To highlight some of our algorithm’s limitations, we design a task specifically to be impossible for unsupervised elicitation. Suppose we really like poems about the sun, so we construct a comparison dataset where all poems that mention the word "sun" are preferred. The only task description we give the LMs is to judge which poem is better, but it is impossible for the LM to know our specific personal preference about poems. In other words, this task is not “salient” to pretrained models, because their understanding of the “poem quality” concept is not related to the sun. To construct the dataset, we use Claude 3.5 Sonnet to generate pairs of poems, and use designed prompts and post-filterings to ensure only one of them mentions “sun”. Experiment results with Llama 70B are shown in Figure 5. As expected, we find ICM performs no better than random guessing. + +4.3 Eliciting Superhuman Capabilities + +After studying unsupervised elicitation on three common NLP datasets, we are further interested in tasks where pretrained models are strongly superhuman. To study this, we explore an author gender prediction task using the Blog Authorship Corpus [27]. 2 + +Using pairs of blog posts ( A and B) from the Blog Authorship Corpus, one written by a male and one by a female, the task is to predict which one is more likely to be written by a male. We use the simple asymmetry logical consistency: A > B contradicts B > A . + +Figure 6: Results on gender prediction. To build human baselines, we recruit 5 annotators to label 1) 48 training examples for prompting and 2) 100 test examples for estimating human performance on the whole test set. Human labels have perfect consistency but bad accuracy (60% on the test set, 53.8% on the training set). As shown in Figure 6, our method matches golden super-vision (80% accuracy), significantly outperforming the estimated human accuracy (60%). In comparison, prompt-ing with weak human labels or commercial post-training all fail to fully leverage pretrained models’ superhuman-level capability. + +4.4 Training an Assistant Chatbot without Supervision + +Figure 7: Accuracy of reward models (left) and pairwise winrates of assistant policy models against the human-supervised baseline (right). We train a Claude 3 Haiku-based reward model, using the Alpaca data or the production data used for training publicly released Claude 3.5 Haiku. Next, we optimize the Claude 3.5 Haiku pretrained model against our reward model to build an assistant policy. + +> 2Our goal is not to improve AI performance at predicting author gender, but rather to study how well this capability is already present in pretrained models. + +7After verifying ICM on standard benchmarks, we investigate whether it can scale to commercial production runs and improve frontier assistant chatbots. Specifically, we aim to train a helpful chat assistant based on the Claude 3.5 Haiku pretrained model, without introducing any human preferences or supervision labels whatsoever. + +Reward Model Training. We use Claude 3 Haiku 3 to generate unsupervised labels. We use the task description “which response is more helpful, harmless, and honest?”. We sample a subset from the production preference dataset for training the publicly released Claude 3.5 Haiku. This subset consists of nearly 400K examples with a 64K token limit. We first use ICM to label 6K examples, train an initial reward model (RM) to label the rest of the data, and then train the final unsupervised RM. We also run the same process on Alpaca to serve as a baseline against this production data. We conduct evaluations on Rewardbench [ 18 ], a widely-used challenging benchmark for RMs. Figure 7 (left) shows the results. First, the human-supervised RM trained on the production data significantly outperforms the RM trained on Alpaca, due to its high-quality human labels and complex examples. Consequently, surpassing the human-supervised RM trained on production data is much harder. Nevertheless, our unsupervised RM still achieves higher accuracy (75.0% v.s. 72.2%). + +Reinforcement Learning with Unsupervised RM. Using both the unsupervised and human-supervised RM, we train two policies via reinforcement learning to create helpful, harmless, and honest assistants. We train both policies on 20,000 RL episodes. We conduct head-to-head compar-isons between two policies: each model’s responses are graded by the RM for training the publicly released Claude 3.5 Sonnet model. As shown in Figure 7 (right), the policy trained with the unsu-pervised RM achieves a 60% win rate. Both these policies lag severely behind the performance of the publicly released Claude 3.5 Haiku, which achieves a much higher 92% win rate against the human-supervised baseline. This is expected because the publicly released Claude 3.5 Haiku is trained for much longer on a much larger dataset with a Claude 3.5 Haiku-based RM. Overall, these experiments suggest that ICM can scale to commercial production runs. + +# 5 Ablations + +Comparing to randomly perturbed labels. Pretrained models may just be robust to label noise on these benchmarks, thus training labels with a certain level of noise could always match the performance of training on golden labels. To rule out this hypothesis, we construct a set of randomly perturbed labels with the same accuracy as our model-generated labels, and conduct ablation studies with Llama pretrained models with many-shot prompting. As shown in Figure 8, our model-generated labels always achieve substantially better performance. We suspect this is because our labels are more aligned with the model’s understanding of correct labels for the task. + +Figure 8: ICM-produced labels outperform equally accurate randomly perturbed labels. + +Evaluating robustness to worst-case initialization. It is possible that our algorithm could collapse under bad initialization (e.g., all initial K labels are wrong), but we coincidentally never encounter such scenarios in Sec. 4 because they happen rarely. We thus investigate ICM’s robustness against different initializations: • Golden : using golden dataset labels. This corresponds to a semi-supervised setting. + +> 3This was an oversight on our part. Ideally, we would use the same model for reward model and policy. + +8• Random : using random labels (our default setting). • Worst : using entirely wrong labels. + +Figure 9: Impact of initialization. Figure 9 showcases results on TruthfulQA with the Llama 8B model. We report the test accuracy using many-shot prompt-ing. Under random initialization, ICM achieves a comparable average accuracy but a slightly higher variance. Even under worst-case initialization, ICM remains robust, experiencing only a moderate performance drop rather than complete failure. This is mainly due to its iterative nature: a few initial bad labels would not degrade the performance significantly, as they can be gradually corrected as the algorithm progresses. + +Figure 10: Impact of logical consistency. + +Ablating logical consistency. The logical consistency term may be of limited value: ICM only introduces simple and general logical consistency that can be applied to many tasks, because determining fine-grained consistency relationships across examples is challenging. Empirically, we observe different impacts of logical consistency across tasks (Figure 10). For example, on TruthfulQA, removing logical consistency only leads to moderately worse results, as the degenerate solution of solely maximizing mutual predictability (i.e. assigning the same label everywhere) happens rarely. In contrast, logical consistency is crucial on Alpaca, since the degenerate solution almost always happens without that. + +# 6 Related Work + +Scaling beyond Human Supervision. Recent work has shown diverse failure modes of post-training LMs with unreliable human supervision. For example, LMs can learn to reward-hack human-designed supervision signals [ 6] or even real humans themselves [ 31 ]. To scale beyond human supervision, one standard method is to use high-quality verifiable rewards. For example, in math, we can match model outputs with existing ground truth solutions [ 15 ]. Unfortunately, such verifiable rewards are unavailable for most tasks. In contrast, our method can provide superhuman-level supervision in broad tasks, even including creating a general helpful, harmless, and honest assistant. + +Evidence of Latent Capabilities in LMs. Recent work shows that pre-trained base models have already learned strong capabilities for downstream tasks, and post-training in fact does not add much. For example, pretrained models can achieve a comparable or even higher pass@ k than their post-trained counterparts when k is large enough, even when post-training is done with verifiable rewards [ 32 ]. Similarly, pretrained and post-trained models perform nearly identically in decoding, while most distribution shifts occur with stylistic tokens such as discourse markers [ 19 ]. When inspecting model latent representations, recent work also finds that LMs encode strong signals of reasoning correctness [ 33 ] or hallucination [ 17 , 13 ]. However, despite prior empirical evidence about LMs’ latent capabilities, they still fail to elicit them effectively. + +Unsupervised Elicitation of LMs. CCS [ 9 ] is one of the most representative works for unsupervised elicitation, which works by solely using simple logical consistency to find latent knowledge. While moderately outperforming the zero-shot prompting baseline, CCS still significantly underperforms supervised approaches. As argued in [ 12 ], CCS, as well as other unsupervised approaches, often cannot find knowledge, because there are many other prominent features that can satisfy logical consistency properties. Our method addresses this challenge by introducing mutual predictability. Several concurrent studies explore unsupervised elicitation by minimizing label entropy [ 34 , 1], differing from our scoring function. Empirically, these studies focus on math or coding domains using specific Qwen pretrained models. In contrast, our work demonstrates for the first time that unsupervised elicitation algorithms can match or exceed human supervision across pretrained models and a variety of crisp and fuzzy tasks — even including training a general-purpose assistant. Unsupervised elicitation can also be thought of as a special case of weak-to-strong generalization [ 8, 16 ]: while they try to use weak human supervision to elicit strong LMs, we seek to ignore the weak human supervision altogether. 97 Discussion + +The role of logical consistency. At first glance, ICM might look like a consistency-based algorithm, and consistency is indeed part of our scoring function (Sec. 2.2). However, as Sec. 5 shows, removing consistency in our scoring function often does not degrade the maximal performance, but increases the variance. Specifically, the algorithm becomes more likely to collapse into degenerate solutions (that have low logical consistency), like assigning the same label to all data points. Therefore, we understand mutual predictability as the most important term that leads to our empirical success. In particular, mutual predictability also likely enforces complex (probabilistic) consistencies, which cannot be easily captured by general axiomatic logical checks. + +Unsupervised elicitation as an alignment method. In practice, when using unsupervised elicitation for alignment, we would still need humans in the loop for various parts of the post-training process. For example, ICM can be directly applied to enhance constitutional AI [ 5 ] for aligning LMs. Specifically, for each human-specified constitution, we can replicate our pipeline in Sec. 4.4: use ICM to label which assistant response follows the constitution more accurately and train an unsupervised reward model, then use reinforcement learning to optimize and align the assistant towards the constitution. Additionally, we still need humans to validate whether the model is interpreting the constitution as intended, for example using scalable oversight techniques [26, 22, 30]. + +Limitations. Our algorithm has two important limitations: (1) As shown in Sec. 4.2, it cannot elicit any concepts or skills unless they are “salient” to the pretrained model. (2) It doesn’t work with long inputs because we need to fit many dataset examples into the model’s effective context window when calculating the scoring function, particularly for the mutual predictability term. + +Conclusion. As LMs advance, they will become capable of doing tasks that humans struggle to evaluate. Therefore, we need new algorithms beyond RLHF to ensure that they still act in accordance with human intent. Our results suggest that unsupervised elicitation is a promising avenue to elicit specific skills from the model without being bounded by the ability of humans. + +# Acknowledgments + +We would like to thank Alec Radford, Akbir Khan, Monte MacDiarmid, Fabien Roger, John Schulman, Lijie Chen, Ruiqi Zhong, and Jessy Lin for their valuable feedback. + +# References + +[1] Shivam Agarwal, Zimin Zhang, Lifan Yuan, Jiawei Han, and Hao Peng. The unreasonable effectiveness of entropy minimization in llm reasoning. arXiv preprint arXiv:2505.15134 , 2025. [2] Owura Asare, Meiyappan Nagappan, and Nirmal Asokan. Is github’s copilot as bad as humans at introducing vulnerabilities in code? Empirical Software Engineering , 28(6):129, 2023. [3] Amanda Askell, Yuntao Bai, Anna Chen, Dawn Drain, Deep Ganguli, Tom Henighan, Andy Jones, Nicholas Joseph, Ben Mann, Nova DasSarma, et al. A general language assistant as a laboratory for alignment. arXiv preprint arXiv:2112.00861 , 2021. [4] Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et al. Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05862 ,2022. [5] Yuntao Bai, Saurav Kadavath, Sandipan Kundu, Amanda Askell, Jackson Kernion, Andy Jones, Anna Chen, Anna Goldie, Azalia Mirhoseini, Cameron McKinnon, et al. Constitutional ai: Harmlessness from ai feedback. arXiv preprint arXiv:2212.08073 , 2022. [6] Bowen Baker, Joost Huizinga, Leo Gao, Zehao Dou, Melody Y Guan, Aleksander Madry, Wojciech Zaremba, Jakub Pachocki, and David Farhi. Monitoring reasoning models for misbehavior and the risks of promoting obfuscation. arXiv preprint arXiv:2503.11926 , 2025. 10 [7] Sébastien Bubeck, Varun Chadrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4, 2023. [8] Collin Burns, Pavel Izmailov, Jan Hendrik Kirchner, Bowen Baker, Leo Gao, Leopold Aschen-brenner, Yining Chen, Adrien Ecoffet, Manas Joglekar, Jan Leike, et al. Weak-to-strong gener-alization: Eliciting strong capabilities with weak supervision. arXiv preprint arXiv:2312.09390 ,2023. [9] Collin Burns, Haotian Ye, Dan Klein, and Jacob Steinhardt. Discovering latent knowledge in language models without supervision. arXiv preprint arXiv:2212.03827 , 2022. [10] Maximin Coavoux, Shashi Narayan, and Shay B Cohen. Privacy-preserving neural representa-tions of text. arXiv preprint arXiv:1808.09408 , 2018. [11] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168 , 2021. [12] Sebastian Farquhar, Vikrant Varma, Zachary Kenton, Johannes Gasteiger, Vladimir Mikulik, and Rohin Shah. Challenges with unsupervised llm knowledge discovery. arXiv preprint arXiv:2312.10029 , 2023. [13] Javier Ferrando, Oscar Obeso, Senthooran Rajamanoharan, and Neel Nanda. Do i know this entity? knowledge awareness and hallucinations in language models. arXiv preprint arXiv:2411.14257 , 2024. [14] Amelia Glaese, Nat McAleese, Maja Tr˛ ebacz, John Aslanides, Vlad Firoiu, Timo Ewalds, Maribeth Rauh, Laura Weidinger, Martin Chadwick, Phoebe Thacker, et al. Improving alignment of dialogue agents via targeted human judgements. arXiv preprint arXiv:2209.14375 , 2022. [15] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948 , 2025. [16] Peter Hase, Mohit Bansal, Peter Clark, and Sarah Wiegreffe. The unreasonable effectiveness of easy training data for hard tasks. arXiv preprint arXiv:2401.06751 , 2024. [17] Saurav Kadavath, Tom Conerly, Amanda Askell, Tom Henighan, Dawn Drain, Ethan Perez, Nicholas Schiefer, Zac Hatfield-Dodds, Nova DasSarma, Eli Tran-Johnson, et al. Language models (mostly) know what they know. arXiv preprint arXiv:2207.05221 , 2022. [18] Nathan Lambert, Valentina Pyatkin, Jacob Morrison, LJ Miranda, Bill Yuchen Lin, Khyathi Chandu, Nouha Dziri, Sachin Kumar, Tom Zick, Yejin Choi, et al. Rewardbench: Evaluating reward models for language modeling. arXiv preprint arXiv:2403.13787 , 2024. [19] Bill Yuchen Lin, Abhilasha Ravichander, Ximing Lu, Nouha Dziri, Melanie Sclar, Khyathi Chandu, Chandra Bhagavatula, and Yejin Choi. The unlocking spell on base llms: Rethinking alignment via in-context learning. arXiv preprint arXiv:2312.01552 , 2023. [20] Stephanie Lin, Jacob Hilton, and Owain Evans. Truthfulqa: Measuring how models mimic human falsehoods. arXiv preprint arXiv:2109.07958 , 2021. [21] Lingjuan Lyu, Xuanli He, and Yitong Li. Differentially private representation for nlp: Formal guarantee and an empirical study on privacy and fairness. arXiv preprint arXiv:2010.01285 ,2020. [22] Nat McAleese, Rai Michael Pokorny, Juan Felipe Ceron Uribe, Evgenia Nitishinskaya, Maja Trebacz, and Jan Leike. Llm critics help catch llm bugs. arXiv preprint arXiv:2407.00215 ,2024. 11 [23] Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang, editors, Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing , pages 11048–11064, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. [24] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems ,35:27730–27744, 2022. [25] Marc Pirlot. General local search methods. European journal of operational research , 92(3):493– 511, 1996. [26] William Saunders, Catherine Yeh, Jeff Wu, Steven Bills, Long Ouyang, Jonathan Ward, and Jan Leike. Self-critiquing models for assisting human evaluators. arXiv preprint arXiv:2206.05802 ,2022. [27] Jonathan Schler, Moshe Koppel, Shlomo Argamon, and James W Pennebaker. Effects of age and gender on blogging. In AAAI spring symposium: Computational approaches to analyzing weblogs , volume 6, pages 199–205, 2006. [28] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Alpaca: A strong, replicable instruction-following model. 2023. [29] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 , 2023. [30] Jiaxin Wen, Ruiqi Zhong, Pei Ke, Zhihong Shao, Hongning Wang, and Minlie Huang. Learning task decomposition to assist humans in competitive programming. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) ,pages 11700–11723, 2024. [31] Jiaxin Wen, Ruiqi Zhong, Akbir Khan, Ethan Perez, Jacob Steinhardt, Minlie Huang, Samuel R Bowman, He He, and Shi Feng. Language models learn to mislead humans via rlhf. arXiv preprint arXiv:2409.12822 , 2024. [32] Yang Yue, Zhiqi Chen, Rui Lu, Andrew Zhao, Zhaokai Wang, Shiji Song, and Gao Huang. Does reinforcement learning really incentivize reasoning capacity in llms beyond the base model? + +arXiv preprint arXiv:2504.13837 , 2025. [33] Anqi Zhang, Yulin Chen, Jane Pan, Chen Zhao, Aurojit Panda, Jinyang Li, and He He. Reasoning models know when they’re right: Probing hidden states for self-verification. arXiv preprint arXiv:2504.05419 , 2025. [34] Xuandong Zhao, Zhewei Kang, Aosong Feng, Sergey Levine, and Dawn Song. Learning to reason without external rewards. arXiv preprint arXiv:2505.19590 , 2025. + +# Appendix A Additional Implementation Details + +A.1 Hyperparameters + +We set the initial temperature T0 = 10 , the final temperature Tmin = 0 .01 , and the cooling rate + +β = 0 .99 . For the coefficient α, we always start with α = 50 . While a large α usually yields labels of higher quality, it may excessively restrict the acceptance criteria, causing the algorithm to frequently reject new labels. Therefore, we may adjust α to a smaller value (20 or 30) based on the search speed on the training data, without reference to any validation data. 12 A.2 Data Statistics + +Table 2 shows the size of train/test splits used for the experiments in Sec. . Table 2: Data size. + +Dataset # Train # Test + +TruthfulQA 2,560 1,000 GSM8K-verification 2,560 2,971 Alpaca 2,048 933 + +# B Compute Costs + +ICM is one form of inference-time scaling. We thus investigate how many forward passes we need to label each datapoint on average. Specifically, we report the statistics based on labeling n = 128 + +datapoints. As shown in Table 3, ICM often requires 2 to 3 forward passes to label each datapoint. Table 3: The average number of forward passes required to label each datapoint with ICM. + +Dataset Avg. # Forward + +TruthfulQA 2.5 GSM8K-verification 3.9 Alpaca 2.0 + +# C Human Annotation + +In Sec. 4.3, we study an author gender prediction task. To establish a human baseline, we recruit 5 annotators from upwork.com , who are all native speakers with extensive experience in reading and writing. Given two blog posts, the annotator is required to review them and select which one is more likely to be written by a male. Overall, we collect 5 human labels for each example. 13 diff --git a/nbs/simple_icm.py b/nbs/simple_icm.py index ae680e3..7159afc 100644 --- a/nbs/simple_icm.py +++ b/nbs/simple_icm.py @@ -14,7 +14,7 @@ import os, sys from dataclasses import dataclass import dotenv from loguru import logger -from openrouter_wrapper.logprobs import openrouter_completion_wlogprobs, get_logprobs_choices # User's wrapper +from openrouter_wrapper.logprobs import openrouter_completion_wlogprobs, get_logprobs_choices, LogprobsNotSupportedError # User's wrapper from typing import List dotenv.load_dotenv() @@ -24,7 +24,6 @@ logger.remove() logger.add(sys.stderr, format="{time} | {level} | {message}", colorize=True) # %% [code] -from dataclasses import field @dataclass class Config: @@ -49,14 +48,14 @@ C = Config( provider_whitelist=[ 'Chutes','Nebius',], ) -C = Config( - model="meta-llama/llama-3.1-70b-instruct", - provider_whitelist=[ 'Cerebras','Nebius',], -) -C = Config( - model="meta-llama/llama-3.1-8b-instruct", - provider_whitelist=[ 'Cerebras','Nebius',], -) +# C = Config( +# model="meta-llama/llama-3.1-70b-instruct", +# provider_whitelist=[ 'Cerebras','Nebius',], +# ) +# C = Config( +# model="meta-llama/llama-3.1-8b-instruct", +# provider_whitelist=[ 'Cerebras','Nebius',], +# ) # quick test of logprob messages = [{'role': 'user', @@ -96,8 +95,8 @@ for idx, item in enumerate(dataset): } data.append(example) -# Limit to small batch for demo (e.g., 64 items) -data = data[:64] +# HACK: Limit to small batch for demo +data = data[:128] logger.info(f"Loaded {len(data)} examples from TruthfulQA-bool") # %% [code] @@ -118,7 +117,10 @@ logger.info("Initialized labels: {}", {k: v['label'] for k, v in demonstrations. # %% [code] # Predict label using in-context prompting (placeholder with OpenAI) +prediction_count = 0 # Global counter for first 3 logs + def predict_label(example_uid, current_demos, config=C): + global prediction_count # Refined few-shot prompt based on get_judge_prompt_fewshot # Sort by consistency_key for relevance, limit to 8 relevant_demos = sorted( @@ -139,15 +141,23 @@ def predict_label(example_uid, current_demos, config=C): try: # Use wrapper for chat completion (messages format) - messages = [{"role": "user", "content": full_prompt}] + messages = [{"role": "user", "content": full_prompt}, {"role": "assistant", "content": "Judgment: "}] response = openrouter_completion_wlogprobs( model_id=config.model, provider_whitelist=config.provider_whitelist, messages=messages, - max_tokens=1, # Just for "1" or "0" + max_tokens=3, # Allow for "1" or "0" temperature=0.0, top_logprobs=5, ) + + # Debug: Log first 3 prompts and responses + if prediction_count < 3: + logger.info(f"Debug Prediction {prediction_count + 1} - UID {example_uid}:") + logger.info(f"Target Prompt: {target_prompt}") + logger.info(f"Response Content: {response['choices'][0]['message']['content']}") + logger.info(f"--- End Debug ---") + prediction_count += 1 # Use wrapper's get_logprobs_choices for score choice_logp, all_logp = get_logprobs_choices(response, ["1", "0"]) score = choice_logp["1"] - choice_logp["0"] @@ -156,24 +166,24 @@ def predict_label(example_uid, current_demos, config=C): # If no logprobs, fallback to text if response['choices'][0]['logprobs'] is None: text = response['choices'][0]['message']['content'].strip() - predicted = 1 if "1" in text or "true" in text.lower() else 0 + + if "0" in text or "false" in text.lower(): + predicted = 0 + elif "1" in text or "true" in text.lower(): + predicted = 1 + else: + predicted = np.nan + logger.error(f"Unclear prediction text: {text}") score = 0.0 logger.warning("No logprobs, using text fallback") logger.info(f"Prediction for UID {example_uid}: {predicted}, score: {score:.2f}") return predicted, float(score) - except LogprobsNotSupportedError as e: - logger.error(f"Logprobs not supported: {e}") - response = openrouter_completion_wlogprobs( - model_id=config.model, - messages=[{"role": "user", "content": full_prompt}], - max_tokens=5, - temperature=0.0, - # No logprobs - ) - text = response['choices'][0]['message']['content'].strip() - predicted = 1 if "1" in text or "true" in text.lower() else 0 - return predicted, 0.0 + # except LogprobsNotSupportedError as e: + # logger.error(f"Logprobs not supported: {e}") + # text = response['choices'][0]['message']['content'].strip() + # predicted = 1 if "1" in text or "true" in text.lower() else 0 + # return predicted, 0.0 except Exception as e: logger.error(f"API error: {e}") return random.choice([0, 1]), 0.0 # Fallback @@ -184,32 +194,101 @@ test_pred, test_score = predict_label(0, current_labeled, C) logger.info(f"Test prediction: {test_pred}, score: {test_score}") # %% [code] -# Simplified inconsistency fix: Ensure per group at most one 'True' for differing keys -def fix_inconsistencies(demos): - # Group by consistency_id - groups = {} - for uid, demo in demos.items(): - cid = demo['consistency_id'] - if cid not in groups: - groups[cid] = [] - groups[cid].append(uid) - +# Enhanced inconsistency fix: Multi-iter LLM proposals for contradictions/implications (inspired by ICM_tools.py) +# FIXME: Current basic heuristic; uses LLM for decision prompts on inconsistent pairs, simulates outcomes, iterates to resolve +async def fix_inconsistencies(demos, config=C, max_iters=3): updated = False - for cid, uids in groups.items(): - trues = [uid for uid in uids if demos[uid]['label'] == 1] - keys = [demos[uid]['consistency_key'] for uid in trues] - if len(trues) > 1 and len(set(keys)) > 1: # Contradiction: multiple trues with different keys - # Simple fix: keep the one with highest score, set others to 0 - scored_trues = [(uid, demos[uid].get('score', 0)) for uid in trues] - best_uid = max(scored_trues, key=lambda x: x[1])[0] - for uid in trues: - if uid != best_uid: - demos[uid]['label'] = 0 - updated = True + for iteration in range(max_iters): + # Group by consistency_id + groups = {} + for uid, demo in demos.items(): + if demo['label'] is not None: + cid = demo['consistency_id'] + if cid not in groups: + groups[cid] = [] + groups[cid].append(uid) + + fixes_made = False + for cid, uids in groups.items(): + labeled_uids = [uid for uid in uids if demos[uid]['label'] is not None] + if len(labeled_uids) < 2: + continue # Need at least two labeled for conflict + + # Find inconsistent pairs (contradiction or implication) + pairs = [] + labels = {uid: demos[uid]['label'] for uid in labeled_uids} + keys = {uid: demos[uid]['consistency_key'] for uid in labeled_uids} + for i, uid1 in enumerate(labeled_uids): + for uid2 in labeled_uids[i+1:]: + label1, label2 = labels[uid1], labels[uid2] + key1, key2 = keys[uid1], keys[uid2] + if key1 != key2 and ((label1 == label2 == 1) or (label1 == label2 == 0 and key1 in ['A>B', 'B>A'])): + pairs.append((uid1, uid2, "contradiction")) + elif key1 == key2 and label1 != label2: + pairs.append((uid1, uid2, "implication")) + + if not pairs: + continue + + # For each pair, use LLM to propose resolution + for uid1, uid2, pair_type in pairs: + claim1 = demos[uid1] + claim2 = demos[uid2] + + # Decision prompt (adapted from get_decision_prompt) + decision_prompt = f"""Resolve inconsistency between two claims: +Claim 1: {claim1['prompt'][:-1]} {claim1['label']} +Claim 2: {claim2['prompt'][:-1]} {claim2['label']} +Type: {pair_type} + +If contradiction, decide which to set True (1) and False (0). +If implication, decide if both True or both False. +Respond with 1 if keep Claim1 True and Claim2 False, or 0 otherwise.""" + + try: + messages = [{"role": "user", "content": decision_prompt}] + response = openrouter_completion_wlogprobs( + model_id=config.model, + provider_whitelist=config.provider_whitelist, + messages=messages, + max_tokens=1, + temperature=0.0, + top_logprobs=5, + ) + choice_logp = get_logprobs_choices(response, ["1", "0"])[0] + decision_score = choice_logp["1"] - choice_logp["0"] + decision = 1 if decision_score > 0 else 0 + except: + decision = 0 # Fallback; prefer Claim1 + + # Apply decision + if pair_type == "contradiction": + if decision == 1: + demos[uid1]['label'] = 1 + demos[uid2]['label'] = 0 + else: + demos[uid1]['label'] = 0 + demos[uid2]['label'] = 1 + else: # implication + if decision == 1: + demos[uid1]['label'] = 1 + demos[uid2]['label'] = 1 + else: + demos[uid1]['label'] = 0 + demos[uid2]['label'] = 0 + + fixes_made = True + updated = True + + if not fixes_made: + break # No more fixes needed in this iteration + return demos, updated -# Test fix -temp_demos, updated = fix_inconsistencies(demonstrations.copy()) +# Test fix (now async, so await) +import asyncio +temp_demos = demonstrations.copy() +temp_demos, updated = asyncio.run(fix_inconsistencies(temp_demos, C)) logger.info(f"Fixed inconsistencies: {updated}") # %% [code] @@ -247,7 +326,7 @@ logger.info("Initial energy: {}", compute_energy(demonstrations)) # %% [code] # Main simulated annealing loop -def run_icm(demonstrations, config=C): +async def run_icm(demonstrations, config=C): energies = [] accuracies = [] current_labeled = {k: v for k, v in demonstrations.items() if v['label'] is not None} @@ -297,7 +376,7 @@ def run_icm(demonstrations, config=C): temp_demos = demonstrations.copy() temp_demos[example_uid]['label'] = new_label temp_demos[example_uid]['score'] = score - temp_demos, _ = fix_inconsistencies(temp_demos) + temp_demos, _ = await fix_inconsistencies(temp_demos, config=config) # Compute new energy new_energy, new_metrics = compute_energy(temp_demos, config) @@ -319,16 +398,20 @@ def run_icm(demonstrations, config=C): if iter % 10 == 0: logger.info("Progress: Labeled {}, Inconsistents: {}", new_metrics['num_labeled'], new_metrics['num_inconsistent']) + # Final async fix if needed + demonstrations, _ = await fix_inconsistencies(demonstrations, config=config) + return demonstrations, energies, accuracies # %% [code] # Run the algorithm -final_demos, energies, accuracies = run_icm(demonstrations, C) +final_demos, energies, accuracies = asyncio.run(run_icm(demonstrations, C)) # Final metrics final_energy, final_metrics = compute_energy(final_demos, C) logger.info("\nFinal Results:") logger.info("Energy: {:.2f}", final_energy) +# TODO show vanilla accuracy here for comparison logger.info("Accuracy vs vanilla: {:.2f}", final_metrics['accuracy']) logger.info("Labeled: {}/{}", final_metrics['num_labeled'], len(data)) logger.info("Inconsistencies: {}", final_metrics['num_inconsistent']) @@ -358,6 +441,7 @@ plt.xlabel('Iteration') plt.ylabel('Accuracy') plt.tight_layout() +plt.savefig("icm_progress.png") plt.show() # %% [markdown] @@ -365,7 +449,9 @@ plt.show() # - [x] Load larger HuggingFace dataset ('Yik/truthfulQA-bool' subset, formatted to messages). # - [x] Refine few-shot prompt from original get_judge_prompt_fewshot. # - [x] Add weighted sampling for inconsistent groups (no longer random). -# - FIXME: Implement async calls for batch efficiency (currently sync). -# - FIXME: -# th logprobs-supported models. - +# - [ ] Implement async batch predictions in predict_label (use asyncio.gather for concurrent API calls, inspired by pipeline.py) +# - [ ] Enhance fix_inconsistencies with multi-iter LLM proposals (use LLM decisions for contradictions/implications, from ICM_tools.py) +# - it's marked as "basic" because it directly applies LLM decisions without deeper simulation or multiple proposals. +# - in the original ICM.py the LLM generates multiple resolution proposals for inconsistencies, simulates their outcomes (e.g., by temporarily applying them and evaluating metrics like energy), and iterates to select the best one. https://github.com/Jiaxin-Wen/Unsupervised-Elicitation/blob/master/src/experiments/ICM.py +# - [ ] Add caching for predictions (dict/file-based, like save_to_cache in pipeline.py) +# - [ ] Expand metrics in compute_energy (add label distributions, detailed inconsistent_num, from ICM.py) diff --git a/src/datatypes/enums.py b/src/datatypes/enums.py deleted file mode 100644 index f8dfaf6..0000000 --- a/src/datatypes/enums.py +++ /dev/null @@ -1,34 +0,0 @@ -__all__ = ["Language", "PromptType"] - -from enum import Enum - - -class Language(Enum): - PYTHON = ("python", "Python") - CPP = ("cpp", "C++") - - def __init__(self, code, text): - self.code = code - self.text = text - - @staticmethod - def from_code(code): - if code == "python": - return Language.PYTHON - elif code == "cpp": - return Language.CPP - else: - raise Exception(f"Unknown code langauge: {code}") - - -class PromptType(Enum): - SOLUTION = "solution_generation" - BLUE_TEAM = "blue_team" - RED_TEAM = "red_team" - EVAL = "eval" - - -class DifficultyEstimationType(Enum): - PROBLEM_ONLY = "problem_only" - PROBLEM_SOLUTION = "problem_solution" - PROBLEM_SOLUTION_EXECUTION = "problem_solution_execution" diff --git a/src/experiments/ICM.py b/src/experiments/ICM.py deleted file mode 100644 index d716693..0000000 --- a/src/experiments/ICM.py +++ /dev/null @@ -1,560 +0,0 @@ -import asyncio -import json -import math -import os -import random -from collections import Counter -from copy import deepcopy -from tqdm import tqdm -import numpy as np -from datasets import load_dataset -import argparse - -from core.llm_api.llm import ModelAPI -from unsupervised_elicitation.utils import setup_environment -from src.experiments.ICM_tools import ( - propose_consistencyfix, - run_consistencyfix, - pick_two_inconsistent_claims, - update_assign_based_on_decision, -) -from src.model_querying.prompt_creation import ( - get_decision_prompt, - get_judge_prompt_fewshot, -) -from src.model_querying.solution_extraction import ( - extract_claim_logprobs, - extract_decision_logprobs, -) -from src.pipeline.pipeline import Pipeline, PipelineConfig -from src.tools.dataloaders import ( - load_assignments, - load_problems_from_json, - load_problems_from_json_ids, -) -from src.tools.path_utils import get_default_results_directory, get_root_directory - - -def calculate_accuracy(train_data, inconsistent_pairs): - train_probs = [] - for i in train_data.values(): - if i["label"] is None: - continue - if i["label"] == 1: - train_probs.append(i["score"]) - else: - train_probs.append(-i["score"]) - if len(train_probs) == 0: - train_prob = 0 - else: - train_prob = np.mean(train_probs) - - return { - "train_accuracy": 0 - if len(train_data) == 0 - else np.mean([i["label"] == i["vanilla_label"] for i in train_data.values()]), - "train_label_distribution": Counter( - [i["vanilla_label"] for i in train_data.values()] - ), - "train_predict_distribution": Counter( - [i["label"] for i in train_data.values()] - ), - "train_prob": train_prob, - "train_size": len(train_data), - "inconsistent_num": len(inconsistent_pairs), - } - - -def update_assign(data): - for key, value in data.items(): - if value["score"] > 0: - value["label"] = 1 - else: - value["label"] = 0 - return data - - -def fix_inconsistency(demonstrations, cur_metric, name, alpha, iter=0, K=20): - backup_metric = deepcopy(cur_metric) - if cur_metric["inconsistent_num"] == 0: - return demonstrations, cur_metric - - cur_pool = {k: v for k, v in demonstrations.items() if v["label"] is not None} - assignment = cur_pool - - best_metric = cur_metric - best_assignment = assignment - best_decision_id = None - for k in range(K): - pipeline = propose_consistencyfix( - args.model, - name=name, - iter=f"{iter}-{k}", - assignment=assignment, - ) - results = asyncio.run(pipeline.run()) - decisions = results["decisions"] - assignment = results["get_assign"] - for decision_id, decision in enumerate(decisions.values()): - tmp_decision_metric_list = [] - tmp_decision_assignment_list = [] - for score_idx, score in enumerate([0, 1]): - tmp_decision = deepcopy(decision) - tmp_decision["score"] = score - tmp_assignment = update_assign_based_on_decision( - deepcopy(assignment), tmp_decision - ) - tmp_pipeline = run_consistencyfix( - model=args.model, - name=name, - iter=f"{iter}-{k}-{decision_id}-{score_idx}", - assignment=tmp_assignment, - ) - tmp_results = asyncio.run(tmp_pipeline.run()) - tmp_metric = tmp_results["evaluate"] - tmp_decision_metric_list.append(tmp_metric) - tmp_decision_assignment_list.append(tmp_assignment) - tmp_best_decision_id = np.argmax( - [get_energy(i, args.alpha) for i in tmp_decision_metric_list] - ) - tmp_assignment = tmp_decision_assignment_list[tmp_best_decision_id] - tmp_metric = tmp_decision_metric_list[tmp_best_decision_id] - - if get_energy(tmp_metric, args.alpha) >= get_energy(best_metric, args.alpha): - best_decision_id = decision_id - best_metric = tmp_metric - best_assignment = tmp_assignment - break - if best_decision_id is None: - break - elif best_metric["inconsistent_num"] == 0: - assignment = best_assignment - break - else: - assignment = best_assignment - - for k in assignment: - demonstrations[k] = assignment[k] - - return demonstrations, best_metric - - -def get_pipeline( - model, - name=None, - use_cache=True, - num_problems=None, - decision_id=None, - iter=None, - assignment=None, -): - pipeline_name = f"iterative-truth-assign-iter-{iter}" - if decision_id is not None: - pipeline_name += f"-{decision_id}" - if name is not None: - pipeline_name += "-" + name - - ROOT_DIR = get_root_directory() - DATA_DIR = ROOT_DIR / "data" - - - pipeline_config = PipelineConfig( - pipeline_name, - anthropic_num_threads=40, - openai_fraction_rate_limit=0.99, - num_problems=num_problems, - use_cache=use_cache, - ) - pipeline = Pipeline(pipeline_config) - - assert assignment is not None - initial_assign = pipeline.add_load_data_step( - "get_assign", load_assignments, assignment - ) - - def add_train_demonstrations(train_data): - copy_data = deepcopy(train_data) - copy_data = {k: v for k, v in copy_data.items() if v["label"] is not None} - keys = list(copy_data.keys()) - values = list(copy_data.values()) - saved_keys = [ - "prompt", - "question", - "choice", - "choice_2", - "consistency_id", - "consistency_key", - "source", - "label", - "vanilla_label", - ] - values = [] - for i in copy_data.values(): - values.append({saved_key: i[saved_key] for saved_key in saved_keys if saved_key in i}) - - for idx, key in enumerate(keys): - tmp_keys, tmp_values = [], [] - for j, (prev_key, prev_value) in enumerate(zip(keys, values)): - if j != idx: - tmp_keys.append(prev_key) - tmp_values.append(prev_value) - - demos = { - prev_key: prev_value - for j, (prev_key, prev_value) in enumerate(zip(tmp_keys, tmp_values)) - } - - sorted_demos = {} - for k, v in demos.items(): - q = v["consistency_id"] - if q not in sorted_demos: - sorted_demos[q] = [] - sorted_demos[q].append((k, v)) - - out_sorted_demos = {} - for group in sorted_demos.values(): - for k, v in group: - out_sorted_demos[k] = v - - copy_data[key]["demonstration"] = out_sorted_demos - - return copy_data - - merged_train_data = pipeline.add_transformation_step( - "add_train_demonstration", - add_train_demonstrations, - dependencies=[initial_assign], - ) - - get_train_preds = pipeline.add_query_step( - "get_train_preds", - model, - get_judge_prompt_fewshot, - extract_claim_logprobs, - dependencies=[merged_train_data], - logprobs=20, - max_tokens=1, - use_cache=use_cache, - ) - - pick_claims = pipeline.add_transformation_step( - "pick_two_inconsistent_claims", - pick_two_inconsistent_claims, - dependencies=[initial_assign], - ) - - eval_preds = pipeline.add_eval_step( - "evaluate", - calculate_accuracy, - dependencies=[get_train_preds, pick_claims], - ) - return pipeline - - -async def predict_assignment(model, example, demonstrations): - demos = [ - v - for k, v in demonstrations.items() - if k != example["uid"] and v["label"] is not None - ] - anthropic_requests = [ - model_api( - model, - get_judge_prompt_fewshot( - example, - demos, - pipeline=False, - ), - logprobs=20, - max_tokens=1, - parse_fn=extract_claim_logprobs, - ) - ] - responses = await asyncio.gather(*anthropic_requests) - score = responses[0][0]["score"] - new_label = score > 0 - return int(new_label) - - -def get_temperature( - iteration, initial_temp, final_temp, decay_rate, schedule="exp" -): - """ - Calculate the temperature for simulated annealing. - - Parameters: - - iteration: Current iteration number. - - initial_temp: Initial temperature. - - decay_rate: Rate at which the temperature decreases. - - Returns: - - Current temperature. - """ - if schedule == "exp": - return max(final_temp, initial_temp * (decay_rate**iteration)) - elif schedule == "log": - return max(final_temp, initial_temp / (1 + 2 * np.log(1 + iteration))) - else: - assert False - - -def get_energy(metric, alpha): - return alpha * metric["train_prob"] - metric["inconsistent_num"] - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--alpha", type=float, default=30) - parser.add_argument("--seed", type=int, default=27565976) - parser.add_argument("--testbed", type=str, default="gsm8k") - parser.add_argument("--model", type=str, default="meta-llama/Llama-3.1-70B") - parser.add_argument("--batch_size", type=int, default=256) - parser.add_argument("--num_seed", type=int, default=8) - parser.add_argument("--K", type=int, default=3000) - parser.add_argument("--consistency_fix_K", type=int, default=10) - parser.add_argument("--decay", type=float, default=0.99) - parser.add_argument("--initial_T", type=float, default=10) - parser.add_argument("--final_T", type=float, default=0.01) - parser.add_argument("--scheduler", type=str, default="log") - args = parser.parse_args() - return args - -def load_data(args): - if args.testbed == "alpaca": - with open(get_root_directory() / "data/train_alpaca.json") as f: - train = json.load(f) - template = """Human: {question} -Response 1: {choice} -Response 2: {choice_2} -Claim: Response 1 is more helpful and harmless than Response 2. -I think this claim is """ - for i in train: - i['source'] = 'alpaca' - i['consistency_key'] = 'A>B' if i['label'] else 'B>A' - i['prompt'] = template.format(question=i['question'], choice=i['choice'].replace("\n", ""), choice_2=i['choice_2'].replace("\n", "")) - - args.GROUP_SIZE = 2 - - elif args.testbed == 'gsm8k': - with open(get_root_directory() / "data/train_gsm8k.json") as f: - train = json.load(f) - template = """Question: {question} -Claim: {answer} -I think this claim is """ - - for i in train: - i['source'] = 'gsm8k' - i['consistency_key'] = i['answer'] - i['prompt'] = template.format(question=i['question'], answer=i['choice']) - - args.GROUP_SIZE = 4 - - elif args.testbed == 'truthfulQA': - with open(get_root_directory() / "data/train_truthfulqa.json") as f: - train = json.load(f) - template = """Question: {question} -Claim: {answer} -I think this claim is """ - - for i in train: - i['source'] = 'truthfulQA' - i['consistency_key'] = 'A' if i['label'] else 'B' - i['prompt'] = template.format(question=i['question'], answer=i['choice']) - args.GROUP_SIZE = 4 - - elif args.testbed == 'truthfulQA-preference': - with open(get_root_directory() / "data/train_truthfulqa_preference.json") as f: - train = json.load(f) - template = """Question: {question} -Answer 1: {choice} -Answer 2: {choice_2} -Claim: Answer 1 is more truthful than Answer 2. -I think this claim is """ - - for i in train: - i['source'] = 'truthfulQA-preference' - i['consistency_key'] = 'A>B' if i['label'] else 'B>A' - i['prompt'] = template.format(question=i['question'], choice=i['choice'], choice_2=i['choice_2']) - args.GROUP_SIZE = 2 - - train_map = {} - for i in train: - if i['consistency_id'] not in train_map: - train_map[i['consistency_id']] = [] - train_map[i['consistency_id']].append(i) - - out = [] - for key in train_map: - out += train_map[key] - train = out - - # sample a batch of batch_size datapoints - fewshot_ids = random.sample( - list(range(len(train)// args.GROUP_SIZE)), args.batch_size // args.GROUP_SIZE - ) - fewshot_ids = [ - i * args.GROUP_SIZE + j for i in fewshot_ids for j in range(args.GROUP_SIZE) - ] - - return train, fewshot_ids - -def initialize(train, fewshot_ids, args): - demonstrations = {} - unlabeled_ids = [] - whole_ids = [] - seed_ids = [] - - random_init_labels = [1] * (args.num_seed // 2) + [0] * (args.num_seed // 2) - random.shuffle(random_init_labels) - - for id, i in enumerate(fewshot_ids): - item = train[i] - item["vanilla_label"] = item["label"] # store dataset labels to measure agreement during the searching process - item["uid"] = id - whole_ids.append(item["uid"]) - if id >= args.num_seed: # set labels to None - item["label"] = None - item["type"] = "predict" - unlabeled_ids.append(item["uid"]) - else: # set random labels - item["type"] = "seed" - item["label"] = random_init_labels[id] - seed_ids.append(item["uid"]) - demonstrations[id] = item - - return demonstrations, unlabeled_ids, whole_ids, seed_ids - - -def main(args): - train, fewshot_ids = load_data(args) - - demonstrations, unlabeled_ids, whole_ids, seed_ids = initialize(train, fewshot_ids, args) - - cur_metric = { - "train_prob": -1e6, - "inconsistent_num": 100000, - "train_accuracy": 1.0, - "train_predict_distribution": {"0": 0, "1": 0}, - "train_label_distribution": {"0": 0, "1": 0}, - } - - print('init random labels = ', Counter([i['label'] for i in demonstrations.values() if i['type'] == 'seed']), 'init label acc = ', np.mean([i['label'] == i['vanilla_label'] for i in demonstrations.values() if i['type'] == 'seed'])) - name = f"{args.testbed}-llama70b-K{args.K}-bc{args.batch_size}_seed{args.seed}-initialsize{args.num_seed}-weighted{args.alpha}-decay{args.decay}-initialT{args.initial_T}-finalT{args.final_T}-scheduler{args.scheduler}" - - iter = 0 - flip_cnt = 0 - example_id = 0 - - for _ in tqdm(range(args.K), desc="searching"): - cur_pool = { - k: v for k, v in demonstrations.items() if v["label"] is not None - } - initial_demos = deepcopy(demonstrations) - if iter == 0: - pipeline = get_pipeline( - args.model, - name=name, - num_problems=None, - iter=iter, - assignment=cur_pool, - ) - results = asyncio.run(pipeline.run()) - cur_metric = results["evaluate"] - - demonstrations, cur_metric = fix_inconsistency( - demonstrations, cur_metric, name, args.alpha, iter=iter, K=args.consistency_fix_K - ) - - cur_pool = { - k: v for k, v in demonstrations.items() if v["label"] is not None - } - - while True: # weighted sampling - candidates_ids = whole_ids - weights = [1 for _ in range(len(candidates_ids))] - for i in candidates_ids: - if i in cur_pool: - same_consistency_group_ids = [j for j in candidates_ids if demonstrations[j]["consistency_id"] == demonstrations[i]["consistency_id"]] - for j in same_consistency_group_ids: - if j not in cur_pool: - weights[j] = 100 - - example_id = random.choices(candidates_ids, k=1, weights=weights)[0] - break - - new_label = asyncio.run( - predict_assignment( - args.model, - demonstrations[example_id], - cur_pool, - ) - ) - - if demonstrations[example_id]["label"] != new_label: - tmp_demonstrations = deepcopy(demonstrations) - tmp_demonstrations[example_id]["label"] = new_label - dummy_metric = { - "train_prob": -1e6, - "inconsistent_num": 100000, - "train_accuracy": 1.0, - "train_predict_distribution": {"0": 0, "1": 0}, - "train_label_distribution": {"0": 0, "1": 0}, - } - - tmp_demonstrations, _ = fix_inconsistency( - tmp_demonstrations, - dummy_metric, - name + "newlabelexplore", - args.alpha, - iter=iter, - K=10, - ) - - tmp_pool = { - k: v - for k, v in tmp_demonstrations.items() - if v["label"] is not None - } - pipeline = get_pipeline( - model=args.model, - name=name, - num_problems=None, - iter=iter, - assignment=tmp_pool, - ) - results = asyncio.run(pipeline.run()) - metric = results["evaluate"] - T = get_temperature( - flip_cnt, args.initial_T, args.final_T, args.decay, schedule=args.scheduler - ) - print(f"iter = {iter}, pool size = {len(cur_pool)}, cur acc = {cur_metric['train_accuracy']}, new acc = {metric['train_accuracy']}, cur score = {get_energy(cur_metric, args.alpha)}, new score = {get_energy(metric, args.alpha)}, cur inconsistent num = {cur_metric['inconsistent_num']}, new inconsistent num = {metric['inconsistent_num']}") - print('cur label distribution = ', Counter([i['label'] for i in demonstrations.values() if i['label'] is not None])) - print('new label distribution = ', Counter([i['label'] for i in tmp_demonstrations.values() if i['label'] is not None])) - - accept_prob = math.exp((get_energy(metric, args.alpha) - get_energy(cur_metric, args.alpha)) / T) - print("accept prob = ", accept_prob) - if random.random() < accept_prob: - print("accept") - demonstrations = tmp_demonstrations - flip_cnt += 1 - cur_metric = metric - with open(f"log_{name}.jsonl", "a") as f: - f.write(json.dumps({ - "iter": iter, - "flip_cnt": flip_cnt, - "acc": cur_metric['train_accuracy'], - "score": get_energy(cur_metric, args.alpha), - }) + "\n") - else: - print("reject") - - print("=" * 100) - iter += 1 - - -if __name__ == "__main__": - setup_environment(logger_level="error") - model_api = ModelAPI(anthropic_num_threads=20, openai_fraction_rate_limit=0.99) - args = get_args() - print("task: ", args.testbed) - random.seed(args.seed) - main(args) diff --git a/src/experiments/ICM_tools.py b/src/experiments/ICM_tools.py deleted file mode 100644 index d238b8f..0000000 --- a/src/experiments/ICM_tools.py +++ /dev/null @@ -1,238 +0,0 @@ -import asyncio -import json -import random -from collections import Counter -from copy import deepcopy - -import numpy as np -from datasets import load_dataset - -from src.model_querying.prompt_creation import ( - get_decision_prompt, - get_judge_prompt_fewshot, -) -from src.model_querying.solution_extraction import ( - extract_claim_logprobs, - extract_decision_logprobs, -) -from src.pipeline.pipeline import Pipeline, PipelineConfig -from src.tools.dataloaders import ( - load_assignments, - load_problems_from_json, - load_problems_from_json_ids, -) -from src.tools.path_utils import get_default_results_directory, get_root_directory - - -def calculate_accuracy(train_data, inconsistent_pairs): - return { - "train_predict_distribution": Counter( - [i["label"] for i in train_data.values()] - ), - "train_label_distribution": Counter( - [i["vanilla_label"] for i in train_data.values()] - ), - "train_accuracy": np.mean( - [i["label"] == i["vanilla_label"] for i in train_data.values()] - ), - "train_prob": np.mean( - [ - i["score"] if i["label"] == 1 else -i["score"] - for i in train_data.values() - ] - ), - "train_size": len(train_data), - "inconsistent_num": len(inconsistent_pairs), - } - - -def update_assign_based_on_decision(data, decision): - if decision["type"] == "contradiction": - if decision["score"] > 0: - data[decision["claim_1"]["uid"]]["label"] = 1 - data[decision["claim_2"]["uid"]]["label"] = 0 - else: - data[decision["claim_1"]["uid"]]["label"] = 0 - data[decision["claim_2"]["uid"]]["label"] = 1 - else: - assert decision["type"] == "implication" - if decision["score"] > 0: - data[decision["claim_1"]["uid"]]["label"] = 1 - data[decision["claim_2"]["uid"]]["label"] = 1 - else: - data[decision["claim_1"]["uid"]]["label"] = 0 - data[decision["claim_2"]["uid"]]["label"] = 0 - return data - - -def pick_two_inconsistent_claims(data): - claims = list(data.values()) - - consistency_groups = {} - for claim in claims: - cid = claim["consistency_id"] - if cid not in consistency_groups: - consistency_groups[cid] = [] - consistency_groups[cid].append(claim) - - inconsistent_pairs = {} - for group in consistency_groups.values(): - labels = [claim["vanilla_label"] for claim in group] - for i in range(len(group)): - for j in range(i + 1, len(group)): - if (group[i]['consistency_key'] != group[j]['consistency_key']) and ( - (group[i]['label'] == group[j]['label'] == 1) or - ( - (group[i]['consistency_key'] in ['A>B', 'B>A']) and (group[i]['label'] == group[j]['label'] == 0) # in comparative tasks, at least one of the two claims is correct - ) - ): - # if (group[i]["vanilla_label"] != group[j]["vanilla_label"]) and ( - # group[i]["label"] == group[j]["label"] - # ): - inconsistent_pairs[len(inconsistent_pairs)] = { - "claim_1": group[i], - "claim_2": group[j], - "consistency_id": group[i]["consistency_id"], - "type": "contradiction", - } - elif (group[i]["consistency_key"] == group[j]["consistency_key"]) and ( - group[i]["label"] != group[j]["label"] - ): - inconsistent_pairs[len(inconsistent_pairs)] = { - "claim_1": group[i], - "claim_2": group[j], - "consistency_id": group[i]["consistency_id"], - "type": "implication", - } - random.shuffle(inconsistent_pairs) - return inconsistent_pairs - - -def propose_consistencyfix( - model, - name=None, - iter=None, - assignment=None, - use_cache=True, -): - pipeline_name = f"propose-consistencyfix-iter-{iter}" - if name is not None: - pipeline_name += "-" + name - pipeline_config = PipelineConfig( - pipeline_name, - anthropic_num_threads=40, - openai_fraction_rate_limit=0.99, - num_problems=None, - use_cache=use_cache, - ) - pipeline = Pipeline(pipeline_config) - - initial_assign = pipeline.add_load_data_step( - "get_assign", load_assignments, assignment - ) - - pick_claims = pipeline.add_transformation_step( - "pick_two_inconsistent_claims", - pick_two_inconsistent_claims, - dependencies=[initial_assign], - ) - - get_decision = pipeline.add_query_step( - "decisions", - model, - get_decision_prompt, - extract_decision_logprobs, - dependencies=[pick_claims], - logprobs=20, - max_tokens=1, - use_cache=use_cache, - ) - return pipeline - -def run_consistencyfix( - model, - name=None, - use_cache=True, - decision_id=None, - decision=None, - iter=None, - assignment=None, -): - pipeline_name = f"consistencyfix-iter-{iter}" - if decision_id is not None: - pipeline_name += f"-{decision_id}" - if name is not None: - pipeline_name += "-" + name - - pipeline_config = PipelineConfig( - pipeline_name, - anthropic_num_threads=40, - openai_fraction_rate_limit=0.99, - num_problems=None, - use_cache=use_cache, - ) - pipeline = Pipeline(pipeline_config) - - assert assignment is not None - initial_assign = pipeline.add_load_data_step( - "get_assign", load_assignments, assignment - ) - - pick_claims = pipeline.add_transformation_step( - "pick_two_inconsistent_claims", - pick_two_inconsistent_claims, - dependencies=[initial_assign], - ) - - def add_train_demonstrations(train_data): - copy_data = deepcopy(train_data) - keys = list(copy_data.keys()) - values = list(copy_data.values()) - saved_keys = [ - "prompt", - "question", - "choice", - "choice_2", - "consistency_id", - "source", - "label", - "vanilla_label", - ] - values = [] - for i in copy_data.values(): - values.append( - {saved_key: i[saved_key] for saved_key in saved_keys if saved_key in i} - ) - - for idx, key in enumerate(keys): - train_data[key]["demonstration"] = { - prev_key: prev_value - for j, (prev_key, prev_value) in enumerate(zip(keys, values)) - if j != idx - } - return train_data - - merged_train_data = pipeline.add_transformation_step( - "add_train_demonstration", - add_train_demonstrations, - dependencies=[initial_assign], - ) - - get_train_preds = pipeline.add_query_step( - "get_train_preds", - model, - get_judge_prompt_fewshot, - extract_claim_logprobs, - dependencies=[merged_train_data], - logprobs=20, - max_tokens=1, - use_cache=use_cache, - ) - - eval_preds = pipeline.add_eval_step( - "evaluate", - calculate_accuracy, - dependencies=[get_train_preds, pick_claims], - ) - - return pipeline \ No newline at end of file diff --git a/src/experiments/plot.py b/src/experiments/plot.py deleted file mode 100644 index f494b2b..0000000 --- a/src/experiments/plot.py +++ /dev/null @@ -1,24 +0,0 @@ -import json -from matplotlib import pyplot as plt -import os - -for file in os.listdir("."): - if file.startswith("log_"): - print(file) - with open(file, "r") as f: - data = [json.loads(i) for i in f] - data = data[:60] - x = list(range(8, 8 + len(data))) - y_score = [i['score'] for i in data] - y_acc = [i['acc'] for i in data] - # plt.subplot(1, 2, 1) - plt.plot(x, y_score, label=f'Acc {max(y_acc):.2f}') - plt.xlabel("# Searched Claims") - plt.ylabel("Score") - # plt.subplot(1, 2, 2) - # plt.plot(x, y_acc) - # plt.xlabel("# Searched Claims") - # plt.ylabel("Accuracy") -plt.legend() -plt.tight_layout() -plt.savefig("log.png", dpi=500) \ No newline at end of file diff --git a/src/llm_api/__init__.py b/src/llm_api/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/llm_api/base_llm.py b/src/llm_api/base_llm.py deleted file mode 100644 index 617d0c7..0000000 --- a/src/llm_api/base_llm.py +++ /dev/null @@ -1,137 +0,0 @@ -import json -import logging -from enum import Enum, auto -from typing import Dict, List, Optional, Protocol - -import attrs -import numpy as np -# from anthropic import AI_PROMPT, HUMAN_PROMPT -from pydantic import BaseModel - -PRINT_COLORS = {"user": "cyan", "system": "magenta", "assistant": "light_green"} -LOGGER = logging.getLogger(__name__) - - -class PromptConfig(BaseModel): - partials: Dict[str, str] = {} - word_limit: Optional[int] = 100 - messages: List[Dict[str, str]] = [] - messages1: List[Dict[str, str]] = [] - messages2: List[Dict[str, str]] = [] - vars: Dict[str, str] = {} - - -class LanguageModelConfig(BaseModel): - model: str - temperature: float = 0.2 - top_p: float = 1.0 - max_tokens: Optional[int] = None - max_words: int = 10000 - min_words: int = 0 - num_candidates_per_completion: int = 1 - timeout: int = 120 - logit_bias: Optional[dict] = None - - -class StopReason(Enum): - MAX_TOKENS = auto() - STOP_SEQUENCE = auto() - TOOL_USE = auto() - - @classmethod - def factory(cls, stop_reason: str) -> "StopReason": - """ - Parses the openai and anthropic stop reasons into a StopReason enum. - """ - if stop_reason in ["max_tokens", "length"]: - return cls.MAX_TOKENS - elif stop_reason in ["stop_sequence", "stop", "end_turn", "eos"]: - return cls.STOP_SEQUENCE - elif stop_reason in ['tool_use', "tool_calls"]: - return cls.TOOL_USE - raise ValueError(f"Invalid stop reason: {stop_reason}") - - def __repr__(self): - return self.name - - -@attrs.frozen() -class LLMResponse: - model_id: str - completion: str - stop_reason: StopReason = attrs.field(converter=StopReason.factory) - cost: float - duration: Optional[float] = None - api_duration: Optional[float] = None - logprobs: Optional[list[dict[str, float]]] = None - - def to_dict(self): - return { - "model_id": self.model_id, - "completion": self.completion, - "stop_reason": self.stop_reason.__repr__(), # Convert to some JSON-serializable format. - "duration": self.duration, - "api_duration": self.api_duration, - "cost": self.cost, - "logprobs": self.logprobs, - } - - -class ModelAPIProtocol(Protocol): - async def __call__( - self, - model_ids: list[str], - prompt, - print_prompt_and_response: bool, - max_attempts: int, - **kwargs, - ) -> list[LLMResponse]: - raise NotImplementedError - - -def messages_to_single_prompt(messages) -> str: - if ( - len(messages) >= 2 - and messages[0]["role"] == "system" - and messages[1]["role"] == "user" - ): - combined_content = messages[0]["content"] + " " + messages[1]["content"] - messages = [{"role": "user", "content": combined_content}] + messages[2:] - prompt = "" - for message in messages: - role = message["role"] - content = message["content"] - tag = AI_PROMPT if role == "assistant" else HUMAN_PROMPT - prompt += f"{tag} {content}" - if tag != AI_PROMPT: - prompt += f"{AI_PROMPT}" - return prompt.strip() - - -def convert_to_prob(log_prob: dict, tokens: list) -> tuple[float, float, float]: - logit1 = log_prob.get(tokens[0], None) - logit2 = log_prob.get(tokens[1], None) - - if logit1 is None: - rating = -100 - LOGGER.warning( - f"Missing token0 {tokens[0]} in log_prob, setting rating to -100.0" - ) - else: - rating = logit1 - - if logit1 is None: - logit1 = -100 - if logit2 is None: - logit2 = -100 - - return rating, logit1, logit2 - - -def add_assistant_message(messages: list[dict], assistant_message: str): - last_role = messages[-1]["role"] - if last_role == "assistant": - messages[-1]["content"] += assistant_message - else: - messages.append({"role": "assistant", "content": assistant_message}) - return messages diff --git a/src/llm_api/llm.py b/src/llm_api/llm.py deleted file mode 100644 index 77aa664..0000000 --- a/src/llm_api/llm.py +++ /dev/null @@ -1,274 +0,0 @@ -import asyncio -import json -import logging -import os -from collections import defaultdict -from itertools import chain -from pathlib import Path -from typing import Callable, Literal, Optional, Union - -import attrs - -from core.llm_api.base_llm import LLMResponse, ModelAPIProtocol -from core.llm_api.openai_llm import ( - BASE_MODELS, - GPT_CHAT_MODELS, - OAIBasePrompt, - OAIChatPrompt, - OpenAIBaseModel, - OpenAIChatModel, -) -from unsupervised_elicitation.utils import load_secrets - -LOGGER = logging.getLogger(__name__) - - -@attrs.define() -class ModelAPI: - openai_fraction_rate_limit: float = attrs.field( - default=0.99, validator=attrs.validators.lt(1) - ) - organization: str = "NYU_ORG" - print_prompt_and_response: bool = False - - _openai_base: OpenAIBaseModel = attrs.field(init=False) - _openai_base_arg: OpenAIBaseModel = attrs.field(init=False) - _openai_chat: OpenAIChatModel = attrs.field(init=False) - - running_cost: float = attrs.field(init=False, default=0) - model_timings: dict[str, list[float]] = attrs.field(init=False, default={}) - model_wait_times: dict[str, list[float]] = attrs.field(init=False, default={}) - - def __attrs_post_init__(self): - secrets = load_secrets() - if self.organization is None: - self.organization = "NYU_ORG" - self._openai_base = OpenAIBaseModel( - frac_rate_limit=self.openai_fraction_rate_limit, - organization=secrets[self.organization], - print_prompt_and_response=self.print_prompt_and_response, - ) - self._openai_base_arg = OpenAIBaseModel( - frac_rate_limit=self.openai_fraction_rate_limit, - organization=secrets["ARG_ORG"], - print_prompt_and_response=self.print_prompt_and_response, - ) - self._openai_chat = OpenAIChatModel( - frac_rate_limit=self.openai_fraction_rate_limit, - organization=secrets[self.organization], - print_prompt_and_response=self.print_prompt_and_response, - ) - Path("./prompt_history").mkdir(exist_ok=True) - - @staticmethod - def _load_from_cache(save_file): - if not os.path.exists(save_file): - return None - else: - with open(save_file) as f: - cache = json.load(f) - return cache - - async def call_single( - self, - model_ids: Union[str, list[str]], - prompt: Union[list[dict[str, str]], str], - max_tokens: int, - print_prompt_and_response: bool = False, - n: int = 1, - max_attempts_per_api_call: int = 10, - num_candidates_per_completion: int = 1, - # is_valid: Callable[[str], bool] = lambda _: True, - parse_fn=lambda _: True, - insufficient_valids_behaviour: Literal[ - "error", "continue", "pad_invalids" - ] = "error", - **kwargs, - ) -> str: - assert n == 1, f"Expected a single response. {n} responses were requested." - responses = await self( - model_ids, - prompt, - max_tokens, - print_prompt_and_response, - n, - max_attempts_per_api_call, - num_candidates_per_completion, - parse_fn, - insufficient_valids_behaviour, - **kwargs, - ) - assert len(responses) == 1, "Expected a single response." - return responses[0].completion - - async def __call__( - self, - model_ids: Union[str, list[str]], - prompt: Union[list[dict[str, str]], str], - print_prompt_and_response: bool = False, - n: int = 1, - max_attempts_per_api_call: int = 50, - num_candidates_per_completion: int = 1, - parse_fn=None, - use_cache: bool = True, - file_sem: asyncio.Semaphore = None, - insufficient_valids_behaviour: Literal[ - "error", "continue", "pad_invalids" - ] = "error", - **kwargs, - ) -> list[LLMResponse]: - """ - Make maximally efficient API requests for the specified model(s) and prompt. - - Args: - model_ids: The model(s) to call. If multiple models are specified, the output will be sampled from the - cheapest model that has capacity. All models must be from the same class (e.g. OpenAI Base, - OpenAI Chat). - prompt: The prompt to send to the model(s). Type should match what's expected by the model(s). - max_tokens: The maximum number of tokens to request from the API - print_prompt_and_response: Whether to print the prompt and response to stdout. - n: The number of completions to request. - max_attempts_per_api_call: Passed to the underlying API call. If the API call fails (e.g. because the - API is overloaded), it will be retried this many times. If still fails, an exception will be raised. - num_candidates_per_completion: How many candidate completions to generate for each desired completion. n*num_candidates_per_completion completions will be generated, then is_valid is applied as a filter, then the remaining completions are returned up to a maximum of n. - parse_fn: post-processing on the generated response - save_path: cache path - use_cache: whether to load from the cache or overwrite it - """ - - assert ( - "max_tokens_to_sample" not in kwargs - ), "max_tokens_to_sample should be passed in as max_tokens." - - if isinstance(model_ids, str): - model_ids = [model_ids] - # # trick to double rate limit for most recent model only - - def model_id_to_class(model_id: str) -> ModelAPIProtocol: - if model_id in ["gpt-4-base", "gpt-3.5-turbo-instruct"]: - return ( - self._openai_base_arg - ) # NYU ARG is only org with access to this model - elif model_id in BASE_MODELS: - return self._openai_base - elif model_id in GPT_CHAT_MODELS or "ft:gpt-3.5-turbo" in model_id: - return self._openai_chat - raise ValueError(f"Invalid model id: {model_id}") - - model_classes = [model_id_to_class(model_id) for model_id in model_ids] - # assert model_classes == self._openai_base - # if model_classes == self._openai_base: - # assert "gpt" not in model_ids[0] - # kwargs['api_base'] = "https://5jfmglryfots6s-8000.proxy.runpod.net/v1" - - if len(set(str(type(x)) for x in model_classes)) != 1: - raise ValueError("All model ids must be of the same type.") - - max_tokens = ( - kwargs.get("max_tokens") if kwargs.get("max_tokens") is not None else 2000 - ) - model_class = model_classes[0] - kwargs["max_tokens"] = max_tokens - # Check if current prompt has already been saved in the save file - # If so, directly return previous result - responses = None - if use_cache and kwargs.get("save_path") is not None: - try: - responses = self._load_from_cache(kwargs.get("save_path")) - except: - logging.error(f"invalid cache data: {kwargs.get('save_path')}") - - # After loading cache, we do not directly return previous results, - # but continue running it through parse_fn and re-save it. - # This is because we may frequently update the parse_fn during development - if responses is None: - num_candidates = num_candidates_per_completion * n - responses = await model_class( - model_ids, - prompt, - print_prompt_and_response, - max_attempts_per_api_call, - n=num_candidates, - **kwargs, - ) - - modified_responses = [] - for response in responses: - self.running_cost += response["response"]["cost"] - if kwargs.get("metadata") is not None: - response["metadata"] = kwargs.get("metadata") - if parse_fn is not None: - response = parse_fn(response) - - self.model_timings.setdefault(response["response"]["model_id"], []).append( - response["response"]["api_duration"] - ) - self.model_wait_times.setdefault( - response["response"]["model_id"], [] - ).append( - response["response"]["duration"] - response["response"]["api_duration"] - ) - modified_responses.append(response) - - if kwargs.get("save_path") is not None: - if file_sem is not None: - async with file_sem: - with open(kwargs.get("save_path"), "w") as f: - json.dump(modified_responses, f, indent=2) - else: - with open(kwargs.get("save_path"), "w") as f: - json.dump(modified_responses, f, indent=2) - return modified_responses[:n] - - def reset_cost(self): - self.running_cost = 0 - - -async def demo(): - model_api = ModelAPI(openai_fraction_rate_limit=0.99) - - oai_chat_messages = [ - [ - {"role": "system", "content": "You are gpt-3.5-turbo."}, - {"role": "user", "content": "who are you!"}, - ], - [ - { - "role": "system", - "content": "You are gpt-4", - }, - {"role": "user", "content": "who are you!"}, - ], - ] - oai_chat_models = ["gpt-3.5-turbo-16k"] - oai_chat_requests = [ - model_api( - oai_chat_models, - prompt=message, - max_tokens=16_000, - n=1, - print_prompt_and_response=False, - ) - for message in oai_chat_messages - ] - answer = await asyncio.gather(*oai_chat_requests) - - for responses in answer: - for i in responses: - print(i.completion) - print("=" * 100) - - costs = defaultdict(int) - for responses in answer: - for response in responses: - costs[response.model_id] += response.cost - - print("-" * 80) - print("Costs:") - for model_id, cost in costs.items(): - print(f"{model_id}: ${cost}") - return answer - - -if __name__ == "__main__": - asyncio.run(demo()) diff --git a/src/llm_api/openai_llm.py b/src/llm_api/openai_llm.py deleted file mode 100644 index d08e039..0000000 --- a/src/llm_api/openai_llm.py +++ /dev/null @@ -1,630 +0,0 @@ -# %% -import asyncio -import json -import logging -import os -import random -import time -from datetime import datetime -from itertools import cycle -from traceback import format_exc -from typing import Optional, Union - -import attrs -import openai -import requests -import tiktoken -from openai.openai_object import OpenAIObject as OpenAICompletion -from tenacity import retry, stop_after_attempt, wait_fixed -from termcolor import cprint - -from core.llm_api.base_llm import ( - PRINT_COLORS, - LLMResponse, - ModelAPIProtocol, -) - -OAIChatPrompt = list[dict[str, str]] -OAIBasePrompt = Union[str, list[str]] -LOGGER = logging.getLogger(__name__) - - -def count_tokens(text: str) -> int: - return len(tiktoken.get_encoding("cl100k_base").encode(text)) - - -def price_per_token(model_id: str) -> tuple[float, float]: - """ - Returns the (input token, output token) price for the given model id. - """ - if model_id == "gpt-4-1106-preview": - prices = 0.01, 0.03 - elif model_id == "gpt-3.5-turbo-1106": - prices = 0.001, 0.002 - elif model_id.startswith("gpt-4"): - prices = 0.03, 0.06 - elif model_id.startswith("gpt-4-32k"): - prices = 0.06, 0.12 - elif model_id.startswith("gpt-3.5-turbo-16k"): - prices = 0.003, 0.004 - elif model_id.startswith("gpt-3.5-turbo"): - prices = 0.0015, 0.002 - elif model_id == "davinci-002": - prices = 0.002, 0.002 - elif model_id == "babbage-002": - prices = 0.0004, 0.0004 - elif model_id == "text-davinci-003" or model_id == "text-davinci-002": - prices = 0.02, 0.02 - elif "ft:gpt-3.5-turbo" in model_id: - prices = 0.012, 0.016 - elif "llama" in model_id.lower() or "mixtral" in model_id.lower(): - prices = 0.0015, 0.002 - elif "o1" in model_id.lower(): - prices = 0.01, 0.03 - else: - prices = 0, 0 - # raise ValueError(f"Invalid model id: {model_id}") - - return tuple(price / 1000 for price in prices) - - -@attrs.define() -class Resource: - """ - A resource that is consumed over time and replenished at a constant rate. - """ - - refresh_rate: float = ( - attrs.field() - ) # How many units of the resource are replenished per minute - value: float = attrs.field(init=False) - total: float = 0 - throughput: float = 0 - last_update_time: float = attrs.field(init=False, factory=time.time) - start_time: float = attrs.field(init=False, factory=time.time) - - def __attrs_post_init__(self): - self.value = self.refresh_rate - - def _replenish(self): - """ - Updates the value of the resource based on the time since the last update. - """ - curr_time = time.time() - self.value = min( - self.refresh_rate, - self.value + (curr_time - self.last_update_time) * self.refresh_rate / 60, - ) - self.last_update_time = curr_time - self.throughput = self.total / (curr_time - self.start_time) * 60 - - def geq(self, amount: float) -> bool: - self._replenish() - return self.value >= amount - - def consume(self, amount: float): - """ - Consumes the given amount of the resource. - """ - assert self.geq( - amount - ), f"Resource does not have enough capacity to consume {amount} units" - self.value -= amount - self.total += amount - - -@attrs.define -class OpenAIModel(ModelAPIProtocol): - frac_rate_limit: float - organization: str - print_prompt_and_response: bool = False - model_ids: set[str] = attrs.field(init=False, default=attrs.Factory(set)) - - # rate limit - token_capacity: dict[str, Resource] = attrs.field( - init=False, default=attrs.Factory(dict) - ) - request_capacity: dict[str, Resource] = attrs.field( - init=False, default=attrs.Factory(dict) - ) - lock_add: asyncio.Lock = attrs.field( - init=False, default=attrs.Factory(asyncio.Lock) - ) - lock_consume: asyncio.Lock = attrs.field( - init=False, default=attrs.Factory(asyncio.Lock) - ) - - @staticmethod - def _assert_valid_id(model_id: str): - raise NotImplementedError - - @staticmethod - async def _get_dummy_response_header(model_id: str): - raise NotImplementedError - - @staticmethod - def _count_prompt_token_capacity(prompt, **kwargs) -> int: - raise NotImplementedError - - async def _make_api_call(self, prompt, model_id, **params) -> list[LLMResponse]: - raise NotImplementedError - - @staticmethod - def _print_prompt_and_response(prompt, responses): - raise NotImplementedError - - @staticmethod - def _create_prompt_history_file(prompt): - filename = f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]}_prompt.txt" - with open(os.path.join("prompt_history", filename), "w") as f: - json_str = json.dumps(prompt, indent=4) - json_str = json_str.replace("\\n", "\n") - f.write(json_str) - - return filename - - @staticmethod - def _add_response_to_prompt_file(prompt_file, responses): - with open(os.path.join("prompt_history", prompt_file), "a") as f: - f.write("\n\n======RESPONSE======\n\n") - json_str = json.dumps( - [response.to_dict() for response in responses], indent=4 - ) - json_str = json_str.replace("\\n", "\n") - f.write(json_str) - - async def add_model_id(self, model_id: str): - self._assert_valid_id(model_id) - if model_id in self.model_ids: - return - - # make dummy request to get token and request capacity - model_metadata = await self._get_dummy_response_header(model_id) - token_capacity = int(model_metadata["x-ratelimit-limit-tokens"]) - request_capacity = int(model_metadata["x-ratelimit-limit-requests"]) - print( - f"got capacities for model {model_id}: {token_capacity}, {request_capacity}" - ) - tokens_consumed = token_capacity - int( - model_metadata["x-ratelimit-remaining-tokens"] - ) - requests_consumed = request_capacity - int( - model_metadata["x-ratelimit-remaining-requests"] - ) - print( - f"consumed capacities for model {model_id}: {tokens_consumed}, {requests_consumed}" - ) - token_cap = token_capacity * self.frac_rate_limit - request_cap = request_capacity * self.frac_rate_limit - if model_id in BASE_MODELS: - token_cap *= ( - 10000 # openai does not track token limit so we can increase it - ) - - print(f"setting cap for model {model_id}: {token_cap}, {request_cap}") - self.model_ids.add(model_id) - token_capacity = Resource(token_cap) - request_capacity = Resource(request_cap) - token_capacity.consume(min(token_cap, tokens_consumed)) - request_capacity.consume(min(request_cap, requests_consumed)) - self.token_capacity[model_id] = token_capacity - self.request_capacity[model_id] = request_capacity - - async def __llama_call__( - self, - model_ids: list[str], - prompt, - print_prompt_and_response: bool, - max_attempts: int, - **kwargs, - ) -> list[LLMResponse]: - kwargs = { - key: value - for key, value in kwargs.items() - if key not in ("save_path", "metadata") - } - - start = time.time() - - async def attempt_api_call(): - api_base_list = [os.environ['LLAMA_API_BASE']] - - kwargs["api_base"] = random.choice(api_base_list) - for model_id in cycle(model_ids): - return await asyncio.wait_for( - self._make_api_call(prompt, model_id, start, **kwargs), - timeout=100, # cloudflare has a 100-second limit for a connection to remain open: https://docs.runpod.io/pods/configuration/expose-ports - ) - - model_ids.sort( - key=lambda model_id: price_per_token(model_id)[0] - ) # Default to cheapest model - model_id = model_ids[0] - prompt = self._process_prompt(prompt) - # prompt_file = self._create_prompt_history_file(prompt) - responses: Optional[list[LLMResponse]] = None - for i in range(max_attempts): - try: - responses = await attempt_api_call() - except Exception as e: - error_info = f"Exception Type: {type(e).__name__}, Error Details: {str(e)}, Traceback: {format_exc()}" - LOGGER.warn( - f"Encountered API error: {error_info}.\nRetrying now. (Attempt {i})" - ) - await asyncio.sleep(1.5**i) - else: - break - - if responses is None: - raise RuntimeError( - f"Failed to get a response from the API after {max_attempts} attempts." - ) - - if self.print_prompt_and_response or print_prompt_and_response: - self._print_prompt_and_response(prompt, responses) - - end = time.time() - LOGGER.debug(f"Completed call to {model_id} in {end - start}s.") - return [ - {"prompt": prompt, "response": response.to_dict()} for response in responses - ] - - async def __call__( - self, - model_ids: list[str], - prompt, - print_prompt_and_response: bool, - max_attempts: int, - **kwargs, - ) -> list[LLMResponse]: - if "gpt" not in model_ids[0]: - return await self.__llama_call__( - model_ids, prompt, print_prompt_and_response, max_attempts, **kwargs - ) - kwargs = { - key: value - for key, value in kwargs.items() - if key not in ("save_path", "metadata") - } - start = time.time() - - async def attempt_api_call(): - for model_id in cycle(model_ids): - async with self.lock_consume: - request_capacity, token_capacity = ( - self.request_capacity[model_id], - self.token_capacity[model_id], - ) - if request_capacity.geq(1) and token_capacity.geq(token_count): - request_capacity.consume(1) - token_capacity.consume(token_count) - else: - await asyncio.sleep(0.01) - continue # Skip this iteration if the condition isn't met - - # Make the API call outside the lock - return await asyncio.wait_for( - self._make_api_call(prompt, model_id, start, **kwargs), timeout=120 - ) - - model_ids.sort( - key=lambda model_id: price_per_token(model_id)[0] - ) # Default to cheapest model - async with self.lock_add: - for model_id in model_ids: - await self.add_model_id(model_id) - if "tool" in prompt[0]: - kwargs["tools"] = prompt[0]["tool"] - if "response_format" in prompt[0]: - kwargs['response_format'] = prompt[0]['response_format'] - prompt = self._process_prompt(prompt) - - token_count = self._count_prompt_token_capacity(prompt, **kwargs) - assert ( - max(self.token_capacity[model_id].refresh_rate for model_id in model_ids) - >= token_count - ), "Prompt is too long for any model to handle." - # prompt_file = self._create_prompt_history_file(prompt) - responses: Optional[list[LLMResponse]] = None - for i in range(max_attempts): - try: - responses = await attempt_api_call() - except Exception as e: - error_info = f"Exception Type: {type(e).__name__}, Error Details: {str(e)}, Traceback: {format_exc()}" - LOGGER.warn( - f"Encountered API error: {error_info}.\nRetrying now. (Attempt {i})" - ) - await asyncio.sleep(1.5**i) - else: - break - - if responses is None: - raise RuntimeError( - f"Failed to get a response from the API after {max_attempts} attempts." - ) - - if self.print_prompt_and_response or print_prompt_and_response: - self._print_prompt_and_response(prompt, responses) - - end = time.time() - LOGGER.debug(f"Completed call to {model_id} in {end - start}s.") - return [ - {"prompt": prompt, "response": response.to_dict()} for response in responses - ] - - -_GPT_4_MODELS = [ - "gpt-4o", - "gpt-4", - "gpt-4-0314", - "gpt-4-0613", - "gpt-4-0125-preview", - "gpt-4-32k", - "gpt-4-32k-0314", - "gpt-4-32k-0613", - "gpt-4-1106-preview", - "gpt-4-turbo", - "gpt-4-turbo-preview", - "gpt-4-turbo-2024-04-09", - "gpt-4o-mini", - "gpt-4o-mini-2024-07-18", - "gpt-4o-2024-11-20", - "o1-preview-2024-09-12", - "o1-mini-2024-09-12", - "deepseek/deepseek-chat", - "meta-llama/llama-3.2-3b-instruct", - "meta-llama/llama-3.2-1b-instruct", - "meta-llama/llama-3.3-70b-instruct", - "mistralai/mistral-7b-instruct", - "meta-llama/llama-3-8b-instruct", - "allenai/olmo-7b-instruct", - "01-ai/yi-large", - "meta-llama/llama-2-70b-chat", - "meta-llama/llama-3.1-8b-instruct", - "meta-llama/llama-3.1-70b-instruct", - "meta-llama/llama-3.1-405b-instruct", - "qwen/qwen-2.5-7b-instruct", - "openai/gpt-4o", - "openchat/openchat-7b", - "ai21/jamba-instruct", - "neversleep/llama-3.1-lumimaid-8b", - "mistralai/mixtral-8x7b-instruct:nitro", - "deepseek/deepseek-r1", - "deepseek/deepseek-r1-distill-llama-70b", - "minimax/minimax-01", - "microsoft/phi-4", - "qwen/qvq-72b-preview", -] -_GPT_TURBO_MODELS = [ - "gpt-3.5-turbo", - "gpt-3.5-turbo-0613", - "gpt-3.5-turbo-16k", - "gpt-3.5-turbo-16k-0613", - "gpt-3.5-turbo-1106", - "gpt-3.5-turbo-0125", -] -GPT_CHAT_MODELS = set(_GPT_4_MODELS + _GPT_TURBO_MODELS) - - -class OpenAIChatModel(OpenAIModel): - def _process_prompt(self, prompt: OAIChatPrompt) -> OAIChatPrompt: - return prompt - - def _assert_valid_id(self, model_id: str): - if "ft:" in model_id: - model_id = model_id.split(":")[1] - assert model_id in GPT_CHAT_MODELS, f"Invalid model id: {model_id}" - - @retry(stop=stop_after_attempt(8), wait=wait_fixed(2)) - async def _get_dummy_response_header(self, model_id: str): - url = "https://api.openai.com/v1/chat/completions" - headers = { - "Content-Type": "application/json", - "Authorization": f"Bearer {openai.api_key}", - "OpenAI-Organization": self.organization, - } - data = { - "model": model_id, - "messages": [{"role": "user", "content": "Say 1"}], - } - response = requests.post(url, headers=headers, json=data) - if "x-ratelimit-limit-tokens" not in response.headers: - raise RuntimeError("Failed to get dummy response header") - return response.headers - - @staticmethod - def _count_prompt_token_capacity(prompt: OAIChatPrompt, **kwargs) -> int: - # The magic formula is: .25 * (total number of characters) + (number of messages) + (max_tokens, or 15 if not specified) - BUFFER = 5 # A bit of buffer for some error margin - MIN_NUM_TOKENS = 20 - - num_tokens = 0 - for message in prompt: - num_tokens += 1 - num_tokens += len(message["content"]) / 4 - - return max( - MIN_NUM_TOKENS, - int(num_tokens + BUFFER) - + kwargs.get("n", 1) * kwargs.get("max_tokens", 15), - ) - - def convert_top_logprobs(self, data): - # Initialize the new structure with only top_logprobs - top_logprobs = [] - - for item in data["content"]: - # Prepare a dictionary for top_logprobs - top_logprob_dict = {} - for top_logprob in item["top_logprobs"]: - top_logprob_dict[top_logprob["token"]] = top_logprob["logprob"] - - top_logprobs.append(top_logprob_dict) - - return top_logprobs - - async def _make_api_call( - self, prompt: OAIChatPrompt, model_id, start_time, **params - ) -> list[LLMResponse]: - LOGGER.debug(f"Making {model_id} call with {self.organization}") - - if params.get("logprobs", None): - params["top_logprobs"] = params["logprobs"] - params["logprobs"] = True - - api_start = time.time() - api_response: OpenAICompletion = await openai.ChatCompletion.acreate(messages=prompt, model=model_id, organization=self.organization, **params) # type: ignore - api_duration = time.time() - api_start - duration = time.time() - start_time - context_token_cost, completion_token_cost = price_per_token(model_id) - context_cost = api_response.usage.prompt_tokens * context_token_cost - completion_cost = api_response.usage.completion_tokens * completion_token_cost - return [ - LLMResponse( - model_id=model_id, - completion=choice.message.content - if "tools" not in params - else choice.message.tool_calls[0]["function"]["arguments"], - stop_reason=choice.finish_reason, - api_duration=api_duration, - duration=duration, - cost=context_cost + completion_cost, - logprobs=self.convert_top_logprobs(choice.logprobs) - if choice.logprobs is not None - else None, - ) - for choice in api_response.choices - ] - - @staticmethod - def _print_prompt_and_response( - prompts: OAIChatPrompt, responses: list[LLMResponse] - ): - for prompt in prompts: - role, text = prompt["role"], prompt["content"] - cprint(f"=={role.upper()}:", "white") - cprint(text, PRINT_COLORS[role]) - for i, response in enumerate(responses): - if len(responses) > 1: - cprint(f"==RESPONSE {i + 1} ({response.model_id}):", "white") - cprint(response.completion, PRINT_COLORS["assistant"], attrs=["bold"]) - print() - - -BASE_MODELS = { - "meta-llama/Llama-3.1-8B", - "meta-llama/Llama-3.1-70B", -} - - -class OpenAIBaseModel(OpenAIModel): - def _process_prompt( - self, prompt: Union[OAIBasePrompt, OAIChatPrompt] - ) -> OAIBasePrompt: - if isinstance(prompt, list) and isinstance(prompt[0], dict): - return messages_to_single_prompt(prompt) - return prompt - - def _assert_valid_id(self, model_id: str): - assert model_id in BASE_MODELS, f"Invalid model id: {model_id}" - - @retry(stop=stop_after_attempt(8), wait=wait_fixed(2)) - async def _get_dummy_response_header(self, model_id: str): - url = "https://api.openai.com/v1/completions" - headers = { - "Content-Type": "application/json", - "Authorization": f"Bearer {openai.api_key}", - "OpenAI-Organization": self.organization, - } - data = {"model": model_id, "prompt": "a", "max_tokens": 1} - response = requests.post(url, headers=headers, json=data) - if "gpt" in model_id and "x-ratelimit-limit-tokens" not in response.headers: - raise RuntimeError("Failed to get dummy response header") - return response.headers - - @staticmethod - def _count_prompt_token_capacity(prompt: OAIBasePrompt, **kwargs) -> int: - max_tokens = kwargs.get("max_tokens", 15) - n = kwargs.get("n", 1) - completion_tokens = n * max_tokens - - tokenizer = tiktoken.get_encoding("cl100k_base") - if isinstance(prompt, str): - prompt_tokens = len(tokenizer.encode(prompt)) - return prompt_tokens + completion_tokens - else: - prompt_tokens = sum(len(tokenizer.encode(p)) for p in prompt) - return prompt_tokens + completion_tokens - - async def _make_api_call( - self, prompt: OAIBasePrompt, model_id, start_time, **params - ) -> list[LLMResponse]: - LOGGER.debug(f"Making {model_id} call with {self.organization}") - api_start = time.time() - api_response: OpenAICompletion = await openai.Completion.acreate(prompt=prompt, model=model_id, organization=self.organization, **params) # type: ignore - api_duration = time.time() - api_start - duration = time.time() - start_time - if "gpt" not in model_id: - return [ - LLMResponse( - model_id=model_id, - completion=choice.text, - stop_reason=choice.finish_reason, - api_duration=api_duration, - duration=duration, - cost=0, - logprobs=choice.logprobs.top_logprobs - if choice.logprobs is not None - else None, - ) - for choice in api_response.choices - ] - else: - context_token_cost, completion_token_cost = price_per_token(model_id) - context_cost = api_response.usage.prompt_tokens * context_token_cost - return [ - LLMResponse( - model_id=model_id, - completion=choice.text, - stop_reason=choice.finish_reason, - api_duration=api_duration, - duration=duration, - cost=context_cost / len(api_response.choices) - + count_tokens(choice.message.content) * completion_token_cost, - logprobs=choice.logprobs.top_logprobs - if choice.logprobs is not None - else None, - ) - for choice in api_response.choices - ] - - @staticmethod - def _print_prompt_and_response(prompt: OAIBasePrompt, responses: list[LLMResponse]): - prompt_list = prompt if isinstance(prompt, list) else [prompt] - responses_per_prompt = len(responses) // len(prompt_list) - responses_list = [ - responses[i : i + responses_per_prompt] - for i in range(0, len(responses), responses_per_prompt) - ] - for i, (prompt, response_list) in enumerate(zip(prompt_list, responses_list)): - if len(prompt_list) > 1: - cprint(f"==PROMPT {i + 1}", "white") - if len(response_list) == 1: - cprint(f"=={response_list[0].model_id}", "white") - cprint(prompt, PRINT_COLORS["user"], end="") - cprint( - response_list[0].completion, - PRINT_COLORS["assistant"], - attrs=["bold"], - ) - else: - cprint(prompt, PRINT_COLORS["user"]) - for j, response in enumerate(response_list): - cprint(f"==RESPONSE {j + 1} ({response.model_id}):", "white") - cprint( - response.completion, PRINT_COLORS["assistant"], attrs=["bold"] - ) - print() - - -# %% diff --git a/src/model_querying/__init__.py b/src/model_querying/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/model_querying/prompt_creation.py b/src/model_querying/prompt_creation.py deleted file mode 100644 index c65b4d1..0000000 --- a/src/model_querying/prompt_creation.py +++ /dev/null @@ -1,45 +0,0 @@ -__all__ = ["get_apps_solution_prompt", "get_two_shot_prompt", "get_backdoor_prompt"] - -import json -import random - -from src.datatypes.enums import Language, PromptType -from src.tools.dataloaders import load_prompts -from src.tools.string_manipulation import format_key_suffix, strip_comments_from_string - -SOLUTION_PROMPTS = load_prompts(PromptType.SOLUTION) -RED_TEAM_PROMPTS = load_prompts(PromptType.RED_TEAM) -BLUE_TEAM_PROMPTS = load_prompts(PromptType.BLUE_TEAM) -EVAL_PROMPTS = load_prompts(PromptType.EVAL) - - -class Prompt: - @staticmethod - def empty(): - return Prompt([]) - - def __init__(self, text, logit_bias=None): - self.text = text - self.logit_bias = logit_bias - - -def get_decision_prompt(example): - prompt = 'dummy prompt' - return Prompt(prompt) - - -def get_judge_prompt_fewshot(example, demonstrations=None, pipeline=True): - if demonstrations is None: - demonstrations = list(example["demonstration"].values()) - prompt = "" - for i in demonstrations: - prompt += i['prompt'] - prompt += "True" if i["label"] else "False" - prompt += "\n\n" - - prompt += example['prompt'] - - if pipeline: - return Prompt(prompt) - else: - return prompt \ No newline at end of file diff --git a/src/model_querying/solution_extraction.py b/src/model_querying/solution_extraction.py deleted file mode 100644 index 8805f61..0000000 --- a/src/model_querying/solution_extraction.py +++ /dev/null @@ -1,55 +0,0 @@ -import json -import logging -import math -from copy import copy - -logger = logging.getLogger(__name__) - - -def get_yes_no(x): - x = x.lower() - y = "true" in x - n = "false" in x - if y == n: - return None - return y - - -def get_yes_no_diff_logprobs(logprobs): - eps = 1e-5 - prob_sums = {False: eps, True: eps} - for k, v in logprobs.items(): - o = get_yes_no(k) - if o is None: - continue - prob_sums[o] += math.exp(v) - - if prob_sums[False] == eps and prob_sums[True] == eps: - return 0 - else: - return math.log(prob_sums[True]) - math.log(prob_sums[False]) - - -def extract_claim_logprobs(response): - response = response.copy() - try: - logprobs = response["response"]["logprobs"][0] - response[f"score"] = get_yes_no_diff_logprobs(logprobs) - except Exception as e: - logger.info( - f"Problem {response['metadata']['uid']}: Error extracting judgment: {repr(e)}" - ) - response["score"] = 0 - return response - -def extract_decision_logprobs(response): - response = response.copy() - try: - logprobs = response["response"]["logprobs"][0] - response[f"score"] = get_yes_no_diff_logprobs(logprobs) - except Exception as e: - logger.info( - f"Problem {response['metadata']['uid']}: Error extracting decision: {repr(e)}" - ) - response["score"] = 0 - return response diff --git a/src/pipeline/README.md b/src/pipeline/README.md deleted file mode 100644 index d371fbc..0000000 --- a/src/pipeline/README.md +++ /dev/null @@ -1,27 +0,0 @@ -How to use the pipeline: -First, outline the graph you would like to execute, including the following: -* Data Loading -* Model queries -* Code Execution Eval -* Transformations -* Monitoring - -Next, convert each of the nodes in that graph into the corresponding helper function: -* add_load_data_step -* add_query_step -* add_code_evaluation_step -* add_transformation_step -* add_monitoring_step - -Each of these takes different parameters that you can see in the method signatures. The important ones to know are these: -LoadData takes either a data-loading function and a location, or it takes a dataset -Queries take a prompt function that they pass the incoming data into to create the associated prompt, and a parse function that they use to parse the LLM response -Code Evals take an executor function that executes all of the code in the Solution objects on the associated test cases. -Transforms take arbitrary functions that they apply to the data as a whole. -Monitoring steps take arbitrary monitoring steps. I may eventually enforce that all pipelines end in one of these because it's really what we care about. - -Finally, put the dependencies of each step into their dependencies parameter. This is how execution order is determined and how data flows between steps. If you rely on more than one step, the data will be passed to ordered args in the same order as the list of dependencies. - -You will also need to include a PipelineConfig parameter that contains metadata around how many concurrents to use and similar. - -Once you have a pipeline definition, call the pipeline.run() function on it to execute the graph. This method returns the Pipeline.Results object back, which holds the output of each step in a dictionary. diff --git a/src/pipeline/pipeline.py b/src/pipeline/pipeline.py deleted file mode 100644 index d24feec..0000000 --- a/src/pipeline/pipeline.py +++ /dev/null @@ -1,314 +0,0 @@ -__all__ = ["PipelineConfig", "Pipeline"] - -import asyncio -import logging -from collections import deque - -from tqdm.auto import tqdm - -from core.llm_api.llm import ModelAPI -from src.datatypes.enums import Language -from src.runners.query_model import QueryConfigBuilder, query_model -from src.tools.dataloaders import read_from_cache, save_to_cache -from src.tools.path_utils import get_root_directory - -logger = logging.getLogger(__name__) - - -def in_notebook(): - try: - from IPython import get_ipython - - if get_ipython() is None or "IPKernelApp" not in get_ipython().config: - return False - except ImportError: - return False - return True - - -class Task: - def __init__(self, name, func, use_cache, dependencies=[]): - self.name = name - self.func = func - self.use_cache = use_cache - self.index = None - self.dependencies = dependencies - self.dependents = [] - self.result = None - for dep in dependencies: - dep.dependents.append(self) - - async def execute(self, results): - if self.result is None: - dep_results = [results[dep.name] for dep in self.dependencies] - if asyncio.iscoroutinefunction(self.func): - self.result = await self.func( - *dep_results, use_cache=self.use_cache, index=self.index - ) - else: - self.result = self.func( - *dep_results, use_cache=self.use_cache, index=self.index - ) - return self.result - - -class PipelineConfig: - def __init__( - self, - name, - anthropic_num_threads=2, - openai_fraction_rate_limit=0.99, - use_cache=True, - language=Language.PYTHON, - num_problems=None, - problem_ids=None, - num_open_files=1000000, - organization="NYU_ORG", - print_prompt_and_response=False, - api_base=None, - ): - self.name = name - self.anthropic_num_threads = anthropic_num_threads - self.openai_fraction_rate_limit = openai_fraction_rate_limit - self.organization = organization - self.print_prompt_and_response = print_prompt_and_response - self.use_cache = use_cache - self.language = language - self.num_problems = num_problems - self.problem_ids = problem_ids - self.num_open_files = num_open_files - self.api_base = api_base - self.play_sound = in_notebook() - - -class Pipeline: - def __init__(self, config): - self.config = config - self.steps = [] - self.step_names = set() - self.results = {} - self.model_api = ModelAPI( - self.config.anthropic_num_threads, - self.config.openai_fraction_rate_limit, - self.config.organization, - self.config.print_prompt_and_response, - ) - self.file_sem = asyncio.BoundedSemaphore(self.config.num_open_files) - self.cost = {"red": 0, "blue": 0} - - def add_load_data_step( - self, name, dataloader_fn, data_location, dependencies=[], use_cache=None - ): - if name in self.step_names: - raise ValueError(f"Step name {name} already exists") - self.step_names.add(name) - - def call(*args, use_cache, index): - return dataloader_fn( - data_location, - num_problems=self.config.num_problems, - problem_ids=self.config.problem_ids, - ) - - task = Task(name, call, use_cache, dependencies) - self.steps.append(task) - return task - - def add_query_step( - self, - name, - model, - prompt_fn, - parse_fn, - dependencies=[], - use_cache=None, - temperature=None, - logprobs=None, - team=None, - max_tokens=4096, - bon=1, - ): - if name in self.step_names: - raise ValueError(f"Step name {name} already exists") - self.step_names.add(name) - - query_config_builder = ( - QueryConfigBuilder() - .with_model_to_test(model) - .with_prompt_fn(lambda x: prompt_fn(x)) - .with_parse_fn(lambda x: parse_fn(x)) - .with_num_problems(self.config.num_problems) - .with_max_tokens(max_tokens) - .with_temperature(temperature) - .with_logprobs(logprobs) - .with_bon(bon) - ) - - async def call(data, use_cache, index): - response_dict = await query_model( - self.model_api, - self.file_sem, - query_config_builder.with_experiment_name( - f"{self.config.name}/{index:02d}-{name}" - ) - .with_use_cache(use_cache) - .with_data(data) - .build(), - ) - self.add_cost_data(team, response_dict) - return response_dict - - step = Task(name, call, use_cache, dependencies) - self.steps.append(step) - return step - - def add_transformation_step( - self, - name, - transformation_fn, - dependencies=[], - use_cache=None, - strong_model=None, - weak_model=None, - read_cache=False, - ): - if name in self.step_names: - raise ValueError(f"Step name {name} already exists") - self.step_names.add(name) - - async def call(*args, use_cache, index): - incoming_problem_ids = set().union(*[arg.keys() for arg in args]) - if use_cache and read_cache: - logger.debug( - f"Reading from cache for transformation: {self.config.name}/{index:02d}-{name}/{strong_model}{'+' if strong_model and weak_model else ''}{weak_model}" - ) - data, cached_problem_ids = read_from_cache( - f"{self.config.name}/{index:02d}-{name}/{strong_model}{'+' if strong_model and weak_model else ''}{weak_model}" - ) - if incoming_problem_ids.issubset(set(cached_problem_ids)): - return {k: v for k, v in data.items() if k in incoming_problem_ids} - - if asyncio.iscoroutinefunction(transformation_fn): - output = await transformation_fn(*args) - else: - output = transformation_fn(*args) - - async with self.file_sem: - save_to_cache( - output, - f"{self.config.name}/{index:02d}-{name}/{strong_model}{'+' if strong_model and weak_model else ''}{weak_model}", - delete_existing=read_cache, - incoming_problem_ids=incoming_problem_ids, - ) - return output - - step = Task(name, call, use_cache, dependencies) - self.steps.append(step) - return step - - def add_eval_step( - self, - name, - eval_fn, - dependencies=[], - strong_model=None, - weak_model=None, - ): - if name in self.step_names: - raise ValueError(f"Step name {name} already exists") - self.step_names.add(name) - - async def call(*args, use_cache, index): - output = eval_fn(*args) - cache_obj = {"summary": output} - async with self.file_sem: - save_to_cache( - cache_obj, - f"{self.config.name}/{index:02d}-{name}/{strong_model}{'+' if strong_model and weak_model else ''}{weak_model}", - ) - return output - - step = Task(name, call, None, dependencies) - self.steps.append(step) - return step - - def topological_sort_tasks(self, tasks): - in_degree = {task: len(task.dependencies) for task in tasks} - - queue = deque([task for task in tasks if in_degree[task] == 0]) - sorted_tasks = [] - task_order = {task: i for i, task in enumerate(tasks)} - - while queue: - task = queue.popleft() - sorted_tasks.append(task) - for dependent in task.dependents: - in_degree[dependent] -= 1 - if in_degree[dependent] == 0: - queue.append(dependent) - queue = deque(sorted(queue, key=lambda t: task_order[t])) - - for i, task in enumerate(sorted_tasks): - task.index = i - return sorted_tasks - - def add_cost_data(self, team, response_dict): - cost = sum( - [response["response"]["cost"] for response in response_dict.values()] - ) - if team is not None: - if team not in self.cost: - self.cost[team] = 0 - self.cost[team] += cost - overall_team = team.split("_")[0] - if overall_team != team: - self.cost[overall_team] += cost - - def set_use_cache(self, tasks): - # This is called after tasks.sort, so we are guaranteed to process all - # dependencies before each task itself. - for task in tasks: - if not self.config.use_cache: - task.use_cache = False - continue - - if task.use_cache is None: - task.use_cache = True - - for dep in task.dependencies: - if not dep.use_cache: - task.use_cache = False - - def speak(self, message): - if self.config.play_sound: - from IPython.display import Javascript, display - - display( - Javascript( - f""" - if(window.speechSynthesis) {{ - var synth = window.speechSynthesis; - synth.speak(new window.SpeechSynthesisUtterance('{message}')); - }} - """ - ) - ) - - async def run(self): - steps = self.topological_sort_tasks(self.steps) - self.set_use_cache(steps) - for task in steps: - logger.info( - f"Starting step {task.index}: {task.name} - Using cache: {task.use_cache}" - ) - try: - self.results[task.name] = await task.execute(self.results) - except Exception as e: - logger.error(f"Error in step {task.index}: {task.name}") - logger.error(e) - self.speak("Pipeline failed sad face") - raise e - logger.info(f"Finished step {task.index}: {task.name}") - self.speak("Jobs done") - logger.info("Run complete!! Nice!! 🚀🚀") - return self.results diff --git a/src/runners/__init__.py b/src/runners/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/runners/evaluate_code.py b/src/runners/evaluate_code.py deleted file mode 100644 index 6c77c5d..0000000 --- a/src/runners/evaluate_code.py +++ /dev/null @@ -1,191 +0,0 @@ -__all__ = [ - "EvalConfig", - "EvalConfigBuilder", - "evaluate_solutions", - "examine_solution", - "print_eval", -] - -import json -import os - -from src.code_evaluation.test_results import Solution - -import src.tools.path_utils as path_utils - -DEFAULT_RESULTS_DIR = path_utils.get_default_results_directory() - - -class EvalConfig: - def __init__( - self, - experiment_name, - model_to_test, - executor_fn, - language, - use_cache=True, - data=None, - dataloader_fn=None, - data_location=None, - ): - self.experiment_name = experiment_name - self.model_to_test = model_to_test - if dataloader_fn is not None: - self.data = dataloader_fn(data_location) - elif isinstance(data, str): - with open(json.load(data), "r") as f: - self.data = json.load(f) - else: - self.data = data - self.executor_fn = executor_fn - self.language = language - self.use_cache = use_cache - - def __str__(self): - pass - - def __repr__(self): - return self.__str__() - - -class EvalConfigBuilder: - def __init__(self): - self.experiment_name = None - self.model_to_test = None - self.executor_fn = None - self.language = None - self.data = None - self.use_cache = None - self.dataloader_fn = None - self.data_location = None - - def with_experiment_name(self, experiment_name): - self.experiment_name = experiment_name - return self - - def with_executor_fn(self, executor_fn): - self.executor_fn = executor_fn - return self - - def with_language(self, language): - self.language = language - return self - - def with_model_to_test(self, model_to_test): - self.model_to_test = model_to_test - return self - - def with_use_cache(self, use_cache): - self.use_cache = use_cache - return self - - def with_data(self, data): - self.data = data - return self - - def with_dataloader_fn(self, dataloader_fn): - self.dataloader_fn = dataloader_fn - return self - - def with_data_location(self, data_location): - self.data_location = data_location - return self - - def build(self): - assert self.experiment_name is not None, "Experiment name must be set" - assert self.executor_fn is not None, "Executor function must be set" - assert self.language is not None, "Language must be set" - assert (self.data is not None) or ( - self.dataloader_fn is not None and self.data_location is not None - ), "Data must be set" - return EvalConfig( - self.experiment_name, - self.model_to_test, - self.executor_fn, - self.language, - self.use_cache, - self.data, - self.dataloader_fn, - self.data_location, - ) - - -def evaluate_solutions(eval_config): - if isinstance(eval_config.data, list): - eval_config.data = { - index: response[0] for index, response in enumerate(eval_config.data) - } - - eval_data = [] - for problem_id, response in eval_config.data.items(): - if response == {}: - item = Solution.no_solution(problem_id) - else: - item = Solution.from_response(problem_id, response, eval_config.language) - eval_data.append(item) - - results_dir = DEFAULT_RESULTS_DIR - save_dir = results_dir / eval_config.experiment_name / eval_config.model_to_test - os.makedirs(save_dir, exist_ok=True) - - # load caching results - if eval_config.use_cache and save_dir is not None: - for idx, solution in enumerate(eval_data): - if solution.correct is not None: # already have runtime eval results - continue - save_path = f"{save_dir}/{solution.question_id}.json" - if os.path.exists(save_path): - eval_data[idx] = Solution.from_cache_file(save_path) - - executor_results = eval_config.executor_fn(eval_data) - - for problem_id, problem_data in eval_config.data.items(): - if problem_id not in executor_results: - executor_results[problem_id] = {} - for field, value in problem_data.items(): - if field not in executor_results[problem_id]: - executor_results[problem_id][field] = value - - # save caching results - if save_dir is not None: - for question_id, solution in executor_results.items(): - save_path = f"{save_dir}/{question_id}.json" - with open(save_path, "w") as f: - json.dump(solution, f, indent=2) - - return executor_results - - -def examine_solution(solutions, index): - print( - f"Difficulty:\n{solutions[index][0]['metadata']['difficulty']}\n----------------------------" - ) - print(solutions[index][0]["metadata"]["question"]) - print(solutions[index][0]["solution"]) - print("test cases:") - for test in solutions[index][0]["metadata"]["test_cases"]: - print(f"{test['input']}{test['output']}") - - -def print_eval(results): - correct_tests = 0 - total_tests = 0 - correct_problems = 0 - total_problems = len(results.keys()) - for problem_id, problem_result in results.items(): - correct_tests_local = sum( - [ - 1 - for test_result in problem_result["test_cases"] - if test_result["correct"] - ] - ) - total_tests_local = len(problem_result["test_cases"]) - # print(f"Problem ID: {problem_id}\nTest Results:\n\tCorrect: {correct_tests_local}\n\tTotal: {total_tests_local}\n\tAccuracy: {(correct_tests_local * 100.)/total_tests_local}\nOverall Correct: {problem_result.correct}") - correct_tests += correct_tests_local - total_tests += total_tests_local - if problem_result["correct"]: - correct_problems += 1 - print( - f"Number of Problems: {total_problems}\nNumber Correct: {correct_problems}\nAccuracy: {(correct_problems * 100.)/total_problems}\nNumber of Tests: {total_tests}\nNumber Correct: {correct_tests}\nAccuracy: {(correct_tests * 100.)/total_tests}" - ) diff --git a/src/runners/query_model.py b/src/runners/query_model.py deleted file mode 100644 index 1ec8118..0000000 --- a/src/runners/query_model.py +++ /dev/null @@ -1,276 +0,0 @@ -__all__ = ["QueryConfig", "QueryConfigBuilder", "query_model"] - -import asyncio -import os - -import tiktoken - -import src.tools.path_utils as path_utils - -ROOT_DIR = path_utils.get_root_directory() -DEFAULT_RESULTS_DIR = path_utils.get_default_results_directory() - - -class QueryConfig: - def __init__( - self, - experiment_name, - model_to_test, - dataloader_fn, - data_location, - data, - prompt_fn, - parse_fn=None, - use_cache=False, - num_problems=None, - max_tokens=4096, - results_dir=None, - temperature=None, - logprobs=None, - bon=1, - ): - assert isinstance(model_to_test, str) - self.experiment_name = experiment_name - self.model_to_test = model_to_test - self.dataloader_fn = dataloader_fn - self.data_location = data_location - self.data = data - self.prompt_fn = prompt_fn - self.use_cache = use_cache - self.parse_fn = parse_fn - self.num_problems = num_problems - self.max_tokens = max_tokens - self.results_dir = results_dir - self.temperature = temperature if temperature is not None else 0.0 - self.logprobs = logprobs - self.bon = bon - - def get_data(self): - assert self.data is not None - return self.data - - def __str__(self): - return ( - f"QueryConfig(" - f"experiment_name={self.experiment_name}, " - f"model_to_test={self.model_to_test}, " - f"dataloader_fn={self.dataloader_fn}, " - f"data_location={self.data_location}, " - f"data={self.data}, " - f"prompt_fn={self.prompt_fn}, " - f"use_cache={self.use_cache}, " - f"num_problems={self.num_problems}, " - f"max_tokens={self.max_tokens}, " - f"results_dir={self.results_dir}, " - f"temperature={self.temperature}, " - f"logprobs={self.logprobs}" - ) - - def __repr__(self): - return self.__str__() - - -class QueryConfigBuilder: - def __init__(self): - self.experiment_name = None - self.model_to_test = None - self.dataloader_fn = None - self.data_location = None - self.data = None - self.prompt_fn = None - self.parse_fn = None - self.use_cache = False - self.num_problems = None - self.max_tokens = 4096 - self.results_dir = None - self.temperature = 0.0 - self.logprobs = None - self.bon = 1 - - def with_experiment_name(self, experiment_name): - self.experiment_name = experiment_name - return self - - def with_bon(self, bon): - self.bon = bon - return self - - def with_model_to_test(self, model_to_test): - assert isinstance(model_to_test, str) - self.model_to_test = model_to_test - return self - - def with_dataloader_fn(self, dataloader_fn): - self.dataloader_fn = dataloader_fn - return self - - def with_data_location(self, data_location): - self.data_location = data_location - return self - - def with_data(self, data): - self.data = data - return self - - def with_prompt_fn(self, prompt_fn): - self.prompt_fn = prompt_fn - return self - - def with_parse_fn(self, parse_fn): - self.parse_fn = parse_fn - return self - - def with_use_cache(self, use_cache): - self.use_cache = use_cache - return self - - def with_num_problems(self, num_problems): - self.num_problems = num_problems - return self - - def with_max_tokens(self, max_tokens): - self.max_tokens = max_tokens - return self - - def with_results_dir(self, results_dir): - self.results_dir = results_dir - return self - - def with_temperature(self, temperature): - self.temperature = temperature - return self - - def with_logprobs(self, logprobs): - self.logprobs = logprobs - if logprobs is not None: - assert "claude" not in self.model_to_test - return self - - def build(self): - assert self.experiment_name is not None, "Experiment name must be set" - assert self.model_to_test is not None, "Model to test must be set" - assert self.prompt_fn is not None, "Prompt function must be set" - assert (self.data is not None) or ( - self.dataloader_fn is not None and self.data_location is not None - ), "Data must be set" - assert (self.data is None) or ( - self.dataloader_fn is None and self.data_location is None - ), "Data and dataloader_fn/data_location cannot both be set" - return QueryConfig( - self.experiment_name, - self.model_to_test, - self.dataloader_fn, - self.data_location, - self.data, - self.prompt_fn, - self.parse_fn, - self.use_cache, - self.num_problems, - self.max_tokens, - self.results_dir, - self.temperature, - self.logprobs, - self.bon, - ) - - -def _get_prompts(problems, prompt_fn): - prompts = {} - for problem_id, problem in problems.items(): - prompt = prompt_fn(problem) - if prompt.text: - prompts[problem_id] = prompt - return prompts - - -def get_save_dir(query_config): - results_dir = ( - ROOT_DIR / query_config.results_dir - if query_config.results_dir is not None - else DEFAULT_RESULTS_DIR - ) - save_dir = results_dir / query_config.experiment_name / query_config.model_to_test - os.makedirs(save_dir, exist_ok=True) - return save_dir - - -def move_data_into_metadata(data): - for data_id, value in data.items(): - filtered_val = { - k: v - for k, v in value.items() - if k not in ("metadata", "prompt", "response") - } - metadata = value.get("metadata", {}) - new_metadata = metadata | filtered_val - data[data_id]["metadata"] = new_metadata - - -def format_response(data, model_responses_map, query_config): - model_responses_flattened = {} - for data_id, response in model_responses_map.items(): - if query_config.bon == 1: - model_responses_flattened[f"{data_id}"] = data[data_id] | response[0] - continue - for resp_id, resp in enumerate(response): - model_responses_flattened[f"{data_id}-{resp_id}"] = data[data_id] | resp - return model_responses_flattened - - -def tokenize_logit_bias(logit_bias, model): - tokenizer = tiktoken.encoding_for_model(model) - tokenized_bias = {} - for k, v in logit_bias.items(): - tokenized = tokenizer.encode(k) - assert len(tokenized) == 1, f"Tokenized bias key {k} is not a single token" - tokenized_bias[tokenized[0]] = v - return tokenized_bias - - -async def query_model(model_api, file_sem, query_config): - data = query_config.get_data() - prompts = _get_prompts(data, query_config.prompt_fn) - save_dir = get_save_dir(query_config) - - move_data_into_metadata(data) - - model_requests = [ - model_api( - query_config.model_to_test, - prompts[data_id].text, - max_tokens=query_config.max_tokens, - temperature=query_config.temperature, - n=query_config.bon, - top_p=1.0, - logprobs=query_config.logprobs, - use_cache=query_config.use_cache, - metadata=data[data_id]["metadata"], - parse_fn=query_config.parse_fn, - save_path=f"{save_dir}/{data_id}.json", - file_sem=file_sem, - **( - { - "logit_bias": tokenize_logit_bias( - prompts[data_id].logit_bias, query_config.model_to_test - ) - } - if prompts[data_id].logit_bias is not None - else {} - ), - ) - for data_id in prompts.keys() - ] - - model_responses = await asyncio.gather(*model_requests) - - # pass through data that wasn't modified by the request - model_responses_map = { - data_id: response for data_id, response in zip(prompts.keys(), model_responses) - } - for key in data.keys(): - if key not in prompts: - model_responses_map[key] = [data[key]] - - response = format_response(data, model_responses_map, query_config) - - return response diff --git a/src/tools/__init__.py b/src/tools/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/tools/dataloaders.py b/src/tools/dataloaders.py deleted file mode 100644 index bd74e9c..0000000 --- a/src/tools/dataloaders.py +++ /dev/null @@ -1,264 +0,0 @@ -__all__ = ["load_prompts", "load_problems", "load_problems_from_json", "load_solutions"] - -import json -import logging -import os - -from .path_utils import get_default_results_directory, get_root_directory - -logger = logging.getLogger(__name__) - -ROOT_DIR = get_root_directory() -DATA_DIR = ROOT_DIR / "data" / "APPS" -PROMPTS_DIR = ROOT_DIR / "src" / "prompts" - - -def get_data_dir(): - return DATA_DIR - - -def load_prompts(prompt_type): - files = [f for f in (PROMPTS_DIR / prompt_type.value).glob("*") if f.is_file()] - prompts = {} - for file in files: - with file.open("r") as f: - prompts[file.name] = f.read() - return prompts - - -def load_problem_subset(subset, require_solutions=False, problem_ids=None): - def load_problems(dir, num_problems=None): - if problem_ids: - problem_dirs = problem_ids - else: - problem_dirs = os.listdir(dir) - problems = {} - added = 0 - for problem_dir in problem_dirs: - problem_path = dir / problem_dir - problem = {} - with (problem_path / "metadata.json").open("r") as f: - problem["metadata"] = json.load(f) - if ( - subset != "ALL" - and problem["metadata"]["difficulty"].lower() != subset.lower() - ): - continue - - if require_solutions and not (problem_path / "solutions.json").exists(): - logger.debug( - f"Skipping problem {problem_dir} because it does not have solutions" - ) - continue - - with (problem_path / "question.txt").open("r") as f: - problem["question"] = f.read() - - problem["uid"] = problem_dir - problems[problem_dir] = problem - added += 1 - if added >= num_problems: - break - - return problems - - return load_problems - - -def load_problems(dir, num_problems=None): - return load_problem_subset("ALL")(dir, num_problems) - - -def load_problems_from_json(path, num_problems=None, problem_ids=None): - problems = {} - try: - with open(path) as f: - data = json.load(f) - except Exception as e: - print('read data error: ', e) - data = path - if num_problems is not None: - data = data[:num_problems] - - for i, item in enumerate(data): - item["uid"] = i - if 'vanilla_label' not in item: - item["vanilla_label"] = item["label"] - problems[f"{i}"] = item - return problems - - -def load_problems_from_json_ids(path, num_problems=None, problem_ids=None): - problems = {} - try: - with open(path) as f: - data = json.load(f) - except: - data = path - - if problem_ids is not None: - data = [data[i] for i in problem_ids] - - for i, item in enumerate(data): - item["uid"] = i - if 'vanilla_label' not in item: - item["vanilla_label"] = item["label"] - # acc.append(item['label'] == item['vanilla_label']) - problems[f"{i}"] = item - - return problems - - -def load_assignments(path, num_problems=None, problem_ids=None): - return path - - -def load_solutions(dir, num_problems=None): - solutions = {} - files = dir.glob("*") - files = [f for f in files if f.name != "incoming_problem_ids.json"] - if num_problems is not None: - files = list(files)[:num_problems] - - for file in files: - with file.open("r") as f: - solution = json.load(f) - # Assume 1 solution per file - if isinstance(solution, list): - solutions[file.stem] = solution[0] - else: - solutions[file.stem] = solution - return solutions - - -def load_multiple_solutions(dir, num_problems=None, problem_ids=None): - solutions = {} - files = dir.glob("*") - files = [f for f in files if f.name != "incoming_problem_ids.json"] - - for file in files: - with file.open("r") as f: - solution = json.load(f) - if "metadata" in solution: - solution.pop("metadata") - if "demonstration" in solution: - solution.pop("demonstration") - solutions[file.stem] = solution - return solutions - -def load_multiple_solutions_w2s(dir, num_problems=None, problem_ids=None): - solutions = {} - files = dir.glob("*") - files = [f for f in files if f.name != "incoming_problem_ids.json"] - - for file in files: - with file.open("r") as f: - solution = json.load(f) - metadata = solution[0]['metadata'] - if "demonstration" in metadata: - metadata.pop("demonstration") - metadata['label'] = solution[0]['score'] > 0 - solutions[file.stem] = metadata - return solutions - - -def save_to_cache(data, name, delete_existing=False, incoming_problem_ids=None): - dir = get_default_results_directory() / name - - # Delete all files in the directory first - if delete_existing and os.path.exists(dir): - for file in os.listdir(dir): - file_path = os.path.join(dir, file) - if os.path.isfile(file_path): - os.unlink(file_path) - - os.makedirs(dir, exist_ok=True) - for k, v in data.items(): - if isinstance(v, list): - to_write = [ - { - key: value - for key, value in item.items() - if key not in ["prompt", "response"] - } - for item in v - ] - else: - to_write = { - key: value - for key, value in v.items() - if key not in ["prompt", "response"] - } - with open(dir / f"{k}.json", "w") as f: - json.dump(to_write, f, indent=4) - - if incoming_problem_ids: - with open(dir / "incoming_problem_ids.json", "w") as f: - json.dump({"problem_ids": list(incoming_problem_ids)}, f, indent=4) - - -def read_from_cache(name): - dir = get_default_results_directory() / name - data = {} - incoming_problem_ids = [] - - for file in dir.glob("*.json"): - if file.name == "incoming_problem_ids.json": - with file.open("r") as f: - incoming_problem_ids = json.load(f).get("problem_ids", []) - else: - with file.open("r") as f: - value = json.load(f) - if not value.get("metadata"): - value["metadata"] = {k: v for k, v in value.items()} - data[file.stem] = value - - return data, incoming_problem_ids - - -def load_ground_truth_solutions(problem_ids): - output = {} - for problem_id in problem_ids: - with open(get_data_dir() / "test" / problem_id / "solutions.json", "r") as f: - solutions = json.load(f) - cleaned_solutions = [] - for solution in solutions: - # Remove unwanted lines from the solution - cleaned_solution = [] - for line in solution.split("\n"): - if ( - not line.strip().startswith("#!") - and " input=" not in line - and "sys.stdin" not in line - ): - cleaned_solution.append(line) - cleaned_solutions.append("\n".join(cleaned_solution).strip()) - output[problem_id] = cleaned_solutions - return output - - -def load_test_case(problem_id): - problem_id = problem_id.split("-")[0] - with open(get_data_dir() / "test" / problem_id / "input_output.json", "r") as f: - data = json.load(f) - return [ - {"input": i, "output": o} for (i, o) in zip(data["inputs"], data["outputs"]) - ] - - -def load_test_cases(problem_ids): - output = {} - for problem_id in problem_ids: - output[problem_id] = load_test_case(problem_id) - return output - - -loaded_test_cases = {} - - -def get_test_cases_for_single_problem(problem_id): - global loaded_test_cases - if problem_id not in loaded_test_cases: - loaded_test_cases[problem_id] = load_test_case(problem_id) - - return loaded_test_cases[problem_id] diff --git a/src/tools/path_utils.py b/src/tools/path_utils.py deleted file mode 100644 index cd08b7b..0000000 --- a/src/tools/path_utils.py +++ /dev/null @@ -1,11 +0,0 @@ -__all__ = ["get_root_directory", "get_default_results_directory"] - -from pathlib import Path - - -def get_root_directory(): - return Path(__file__).parent.parent.parent - - -def get_default_results_directory(): - return get_root_directory() / "results" diff --git a/src/tools/printer.py b/src/tools/printer.py deleted file mode 100644 index c767db9..0000000 --- a/src/tools/printer.py +++ /dev/null @@ -1,134 +0,0 @@ -import json -import os -import subprocess -from pathlib import Path - -from anytree import Node, RenderTree -from anytree.exporter import DotExporter - -from src.tools.path_utils import get_default_results_directory - - -def print_experiment_log(experiment_name, strong_model, weak_model, problem_number): - results_dir = get_default_results_directory() - experiment_dir = results_dir / experiment_name - - # Get all step directories - step_dirs = [d for d in experiment_dir.iterdir() if d.is_dir()] - - # Sort step directories by step number and exclude "merged_results" - step_dirs = [d for d in step_dirs if d.name != "merged_results"] - step_dirs.sort(key=lambda x: int(x.name.split("-")[0])) - - for step_dir in step_dirs: - # Check for both strong and weak model directories - for model in [strong_model, weak_model, f"{strong_model}+{weak_model}"]: - model_dir = step_dir / model - if not model_dir.exists(): - continue - - ignore_keys = ["metadata", "prompt", "response"] - if model == f"{strong_model}+{weak_model}": - ignore_keys.extend(["question", "test_cases", "uid"]) - - # Find matching problem files - problem_files = list(model_dir.glob(f"{problem_number}*.json")) - - for problem_file in problem_files: - with open(problem_file, "r") as f: - data = json.load(f) - - print(f"Step: {step_dir.name}") - print(f"Model: {model}") - print(f"Problem: {problem_file.stem}") - if not isinstance(data, list): - data = [data] - - print("\nPrompt:") - prompt_array = data[0].get("prompt") - if prompt_array is None: - print("No prompt available") - else: - for text in prompt_array: - print(f"Role: {text['role']}") - print(f"Content: {text['content']}") - - for response in data: - print("\nResponse:") - print( - response.get("response", {}).get( - "completion", "No response available" - ) - ) - print("\nOther Fields:") - for key, value in response.items(): - if key not in ignore_keys: - print(f"{key}: {value}") - print("\n" + "=" * 50 + "\n") - - -def show_pipeline_graph(pipeline): - import matplotlib.pyplot as plt - import networkx as nx - - # Create a directed graph - G = nx.DiGraph() - - # Add nodes and edges - for task in pipeline.steps: - G.add_node(task.name) - for dep in task.dependencies: - G.add_edge(dep.name, task.name) - - # Print the graph structure - print("Pipeline Dependency Graph:") - for node in nx.topological_sort(G): - predecessors = list(G.predecessors(node)) - successors = list(G.successors(node)) - print(f"{node}:") - if predecessors: - print(f" Parents: {', '.join(predecessors)}") - if successors: - print(f" Children: {', '.join(successors)}") - - # Generate a DOT file for visualization - output_dir = get_default_results_directory() / pipeline.config.name - output_dir.mkdir(parents=True, exist_ok=True) - dot_file = output_dir / "pipeline_graph.dot" - png_file = output_dir / "pipeline_graph.png" - - nx.drawing.nx_pydot.write_dot(G, str(dot_file)) - print(f"DOT file generated at: {dot_file}") - - # Generate PNG file using Graphviz - try: - subprocess.run(["dot", "-Tpng", str(dot_file), "-o", str(png_file)], check=True) - print(f"PNG file generated at: {png_file}") - except subprocess.CalledProcessError: - print( - "Error: Failed to generate PNG. Make sure Graphviz is installed and accessible in your PATH." - ) - except FileNotFoundError: - print( - "Error: Graphviz not found. Please install Graphviz to generate PNG files." - ) - - # Optionally, you can also use matplotlib to visualize the graph - plt.figure(figsize=(12, 8)) - pos = nx.spring_layout(G) - nx.draw( - G, - pos, - with_labels=True, - node_color="lightblue", - node_size=2000, - font_size=8, - arrows=True, - ) - plt.title("Pipeline Dependency Graph") - plt.axis("off") - plt.tight_layout() - plt.savefig(str(output_dir / "pipeline_graph_matplotlib.png")) - print( - f"Matplotlib graph generated at: {output_dir / 'pipeline_graph_matplotlib.png'}" - ) diff --git a/src/tools/string_manipulation.py b/src/tools/string_manipulation.py deleted file mode 100644 index b4ebfd3..0000000 --- a/src/tools/string_manipulation.py +++ /dev/null @@ -1,62 +0,0 @@ -import re - - -def format_key_suffix(key_suffix): - if key_suffix: - if key_suffix[0] != "_": - key_suffix = f"_{key_suffix}" - else: - key_suffix = "" - return key_suffix - - -COMMENTS_REGEX = r""" -^ # Begin of line. -(?: - # A) Capturing group n°1: Full-line comment followed by empty lines. - ( - [ \t]* # Optional spaces or tabs. - \#[^\r\n]*\r?\n # The comment and the new line. - (?:[ \t]*\r?\n)* # Optional empty lines (perhaps with spaces/tabs). - ) -| - # B) Statement and optional comment at the end. - (?: - ( # Capturing group n°2 : The statement - (?: - # Multi-line strings with \"\"\" or '''. - # Capturing group n°3 : The triple quotes. - (['\"]{3})[\s\S]*?\3 - | - # Double-quoted string "It's ok". - \"(?: \\. | [^\"] )*\" - | - # Single-quoted string 'I\'ll say "Hello!"'. - '(?: \\. | [^'] )*' - | - # Any chars, except spaces, hashtag, quotes and new lines. - [^ \t#\"'\r\n]+ - | - # Horizontal spaces, but not followed by a comment, because - # we want the spaces in front of the comment to be matched - # together with the optional comment we want to get rid of. - [ \t]+(?![ \t]*\#) - )+ - ) - # Capturing group n°4: An optional comment at the end of a statement. - ( - [ \t]*\#[^\r\n]* - )? - )+ -) -""" - - -def strip_comments_from_string(string): - return re.sub( - COMMENTS_REGEX, - "\\2", - string, - 0, - re.MULTILINE | re.VERBOSE | re.UNICODE, - ) diff --git a/src/utils.py b/src/utils.py deleted file mode 100644 index 726096e..0000000 --- a/src/utils.py +++ /dev/null @@ -1,239 +0,0 @@ -import asyncio -import json -import logging -import os -import time -from functools import lru_cache, wraps -from pathlib import Path -from typing import Callable - -import matplotlib.pyplot as plt -import numpy as np -import openai -import pandas as pd -import replicate -import typer -import yaml -from tenacity import retry, retry_if_result, stop_after_attempt - -typer.main.get_command_name = lambda name: name - -LOGGER = logging.getLogger(__name__) -SEPARATOR = "---------------------------------------------\n\n" -SEPARATOR_CONVERSATIONAL_TURNS = "=============================================\n\n" -PROMPT_HISTORY = "prompt_history" -SECRETS_FILE_PATH = Path(__file__).parent.parent / "SECRETS" - -LOGGING_LEVELS = { - "critical": logging.CRITICAL, - "error": logging.ERROR, - "warning": logging.WARNING, - "info": logging.INFO, - "debug": logging.DEBUG, -} - - -def setup_environment( - anthropic_tag: str = "ANTHROPIC_API_KEY", - logger_level: str = "info", - openai_tag: str = "API_KEY", - mistral_tag: str = "MISTRAL_API_KEY", - replicate_tag: str = "REPLICATE_API_KEY", - organization: str = None, -): - setup_logging(logger_level) - load_secrets( - SECRETS_FILE_PATH, - anthropic_tag, - logger_level, - openai_tag, - mistral_tag, - replicate_tag, - organization, - ) - - -def setup_logging(level_str): - level = LOGGING_LEVELS.get( - level_str.lower(), logging.INFO - ) # default to INFO if level_str is not found - logging.basicConfig( - level=level, - format="%(asctime)s [%(levelname)s] (%(name)s) %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", - ) - - root_logger = logging.getLogger() - root_logger.setLevel(level) - - # Disable logging from noisy libraries - logging.getLogger("openai").setLevel(logging.CRITICAL) - logging.getLogger("httpx").setLevel(logging.CRITICAL) - logging.getLogger("matplotlib").setLevel(logging.CRITICAL) - logging.getLogger("anthropic").setLevel(logging.CRITICAL) - logging.getLogger("httpcore").setLevel(logging.CRITICAL) - logging.getLogger("urllib3").setLevel(logging.CRITICAL) - LOGGER.info(f"Logging level set to {level_str}") - - -def load_secrets( - file_path=SECRETS_FILE_PATH, - anthropic_tag: str = "ANTHROPIC_API_KEY", - logger_level: str = "info", - openai_tag: str = "API_KEY", - mistral_tag: str = "MISTRAL_API_KEY", - replicate_tag: str = "REPLICATE_API_KEY", - organization: str = None, -): - secrets = {} - with open(file_path) as f: - for line in f: - key, value = line.strip().split("=", 1) - secrets[key] = value - - openai.api_key = secrets[openai_tag] - os.environ['LLAMA_API_BASE'] = secrets['LLAMA_API_BASE'] - # replicate.api_token = secrets[replicate_tag] - # os.environ["ANTHROPIC_API_KEY"] = secrets[anthropic_tag] - # os.environ["MISTRAL_API_KEY"] = secrets[mistral_tag] - # os.environ["REPLICATE_API_KEY"] = secrets[replicate_tag] - - if organization is not None: - openai.organization = secrets[organization] - if secrets.get("API_BASE") is not None: - openai.api_base = secrets['API_BASE'] - return secrets - - -def load_yaml(file_path): - with open(file_path) as f: - content = yaml.safe_load(f) - return content - - -@lru_cache(maxsize=8) -def load_yaml_cached(file_path): - with open(file_path) as f: - content = yaml.safe_load(f) - return content - - -def save_yaml(file_path, data): - with open(file_path, "w") as f: - yaml.dump(data, f) - - -def load_jsonl(file_path): - data = [] - with open(file_path, "r") as f: - for line in f: - json_obj = json.loads(line) - data.append(json_obj) - return data - - -def save_jsonl(file_path, data): - with open(file_path, "w") as f: - for line in data: - json.dump(line, f) - f.write("\n") - - -def delete_old_prompt_files( - path: str = PROMPT_HISTORY, max_age_minutes: int = 60, keep_recent: int = 50 -): - """ - Delete all files in the folder that: - - Are more than max_age_minutes old - - AND are not one of the keep_recent most recent files - """ - if not os.path.exists(path): - return - - # Get all files in the folder with their full paths and creation times - files = [ - { - "path": os.path.join(path, filename), - "ctime": os.path.getctime(os.path.join(path, filename)), - } - for filename in os.listdir(path) - if os.path.isfile(os.path.join(path, filename)) - ] - - # Sort files by creation time - files.sort(key=lambda f: f["ctime"], reverse=True) - - # Current time in seconds since epoch - now = time.time() - - deleted_count = 0 - for index, file_info in enumerate(files): - # File age in minutes - age_minutes = (now - file_info["ctime"]) / 60 - - # If file is older than x_minutes and is not one of the y_most_recent files, delete it - if age_minutes > max_age_minutes and index >= keep_recent: - os.remove(file_info["path"]) - deleted_count += 1 - - if deleted_count > 0: - print(f"Deleted {deleted_count} old prompt files") - - -def typer_async(f): - @wraps(f) - def wrapper(*args, **kwargs): - try: - loop = asyncio.get_running_loop() - except RuntimeError: # No event loop running - loop = None - - if loop is None: - return asyncio.run(f(*args, **kwargs)) - else: - return f(*args, **kwargs) # Return coroutine to be awaited - - return wrapper - - -@retry( - stop=stop_after_attempt(16), - retry=retry_if_result(lambda result: result is not True), -) -def function_with_retry(function, *args, **kwargs): - return function(*args, **kwargs) - - -@retry( - stop=stop_after_attempt(16), - retry=retry_if_result(lambda result: result is not True), -) -async def async_function_with_retry(function, *args, **kwargs): - return await function(*args, **kwargs) - - -def log_model_timings(api_handler, save_location="./model_timings.png"): - if len(api_handler.model_timings) > 0: - plt.figure(figsize=(10, 6)) - for model in api_handler.model_timings: - timings = np.array(api_handler.model_timings[model]) - wait_times = np.array(api_handler.model_wait_times[model]) - LOGGER.info( - f"{model}: response {timings.mean():.3f}, waiting {wait_times.mean():.3f} (max {wait_times.max():.3f}, min {wait_times.min():.3f})" - ) - plt.plot( - timings, label=f"{model} - Response Time", linestyle="-", linewidth=2 - ) - plt.plot( - wait_times, label=f"{model} - Waiting Time", linestyle="--", linewidth=2 - ) - plt.legend() - plt.title("Model Performance: Response and Waiting Times") - plt.xlabel("Sample Number") - plt.ylabel("Time (seconds)") - plt.savefig(save_location, dpi=300) - plt.close() - - -def softmax(x): - return np.exp(x) / np.sum(np.exp(x), axis=0)