Ensembling and ensemble learning methods

import openai
import json
import os
import time  
import pandas as pd
from human_eval.data import write_jsonl, read_problems
from openai import OpenAI
from azure.ai.inference import ChatCompletionsClient
from azure.ai.inference.models import SystemMessage, UserMessage
from azure.core.credentials import AzureKeyCredential
import weave; weave.init("human_eval_moa_gen")


# Initialize LLaMA-based client
llama_client = ChatCompletionsClient(
    temperature=0.0,
    endpoint="https://Meta-Llama-3-1-70B-Instruct-mzte.eastus2.models.ai.azure.com",
    credential=AzureKeyCredential("your key")
)

# Initialize OpenAI client
openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def generate_response_from_model(prompt, model_name):
    if model_name == 'llama':
        response = llama_client.complete(
            messages=[
                SystemMessage(content="You are a helpful assistant."),
                UserMessage(content=prompt)
            ]
        )
        try:
            return response.choices[0].message.content.strip()
        except Exception as e:
            print(f"Failed to get response: {e}")
            return None
    elif model_name == 'openai':
        model = "gpt-4o-mini"
        response = openai_client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "You are an expert Python programmer."},
                {"role": "user", "content": prompt}
            ],
            temperature=0
        )
        return response.choices[0].message.content.strip()
    else:
        raise ValueError(f"Unknown model name: {model_name}")

def save_progress_to_csv(data, filename='progress.csv'):
    df = pd.DataFrame(data)
    if not os.path.isfile(filename):
        df.to_csv(filename, index=False)
    else:
        df.to_csv(filename, mode='a', header=False, index=False)

def load_progress_from_csv(filename='progress.csv'):
    if os.path.isfile(filename):
        return pd.read_csv(filename).to_dict('records')
    return []

def process_problem(task_id, problem, completed_tasks):
    if task_id in completed_tasks:
        print(f"Task {task_id} already processed. Skipping.")
        return None

    prompt = problem["prompt"]
    print(f"Processing task: {task_id}")

    models = ['llama', 'openai']

    responses = []
    for model_name in models:
        full_prompt = (
            f"You are an expert Python programmer. Given the following task and approach, "
            f"NOTE: YOU MUST PUT ALL IMPORTS INSIDE THE FUNCTION, AS ALL CODE OUTSIDE OF SINGLE FUNCTION WILL BE IGNORED."
            f"write the full complete function. "
            f"Handle all of the test cases and reasonable set of edge cases. DO NOT USE VALUEERRORS IN THE CODE. THIS WILL RUIN THE EVAL. JUST HANDLE THE EDGE CASES GRACEFULLY"
            f"ALL CODE MUST BE CONTAINED WITHIN THE SINGLE FUNCTION - EVEN IMPORTS ETC. YOU CAN DEFINE MORE FUNCTIONS BUT THEY MUST BE CONTAINED INSIDE THE CORE FUNCTION DESCRIBED IN THE PROMPT"            
            f"YOU CANNOT USE EXTERNAL LIBRARIES LIKE NUMPY ETC. -- USE STANDARD LIBRARIES ONLY LIKE MATH, from typing import List, Tuple, from itertools import combinations, from collections import Counter, from functools import reduce, import random, import re etc etc."
            f"NOTE: YOU MUST PUT ALL IMPORTS INSIDE THE FUNCTION, AS ALL CODE OUTSIDE OF SINGLE FUNCTION WILL BE IGNORED."
            f"DO NOT ADD ANY EXAMPLE USES. JUST THE FUNCTION!!!!!!!!"
            f"BE CAUTIOUS TO AVOID ' is not defined' errors too"
            f"END CODE WITH '```'. NOTE YOU ABSOLUTELY MUST END THE CODE WITH '```' OR ELSE "
            f"THE CODE WILL NOT BE INTERPRETED!!!!\n\n"
            f"Task:\n{prompt}"
        )
        response = generate_response_from_model(full_prompt, model_name)
        responses.append({'model': model_name, 'response': response})

        # Save individual model responses to JSONL files
        save_individual_response_to_jsonl(task_id, model_name, response)

    print(responses)

    critiques = []
    for model_name in models:
        critique_prompt = (
            f"As an expert Python programmer, analyze the following code for any possible bugs, "
            f"that could prevent test cases from passing or potential (reasonable) edge cases that could make the code invalid.\n\n"
            f"Task:\n{prompt}\n\n"
            f"Code:\n{responses}\n"
        )
        critique = generate_response_from_model(critique_prompt, model_name)
        critiques.append({'model': model_name, 'critique': critique})

    print(critiques)

    final_response = generate_final_response(prompt, responses, critiques)

    print("#" * 50)
    print(final_response)
    print("#" * 50)

    task_result = {'task_id': task_id, 'completion': final_response}
    save_progress_to_csv([task_result])  # Save the result after processing

    return task_result

def save_individual_response_to_jsonl(task_id, model_name, response, output_dir='model_responses'):
    os.makedirs(output_dir, exist_ok=True)
    filename = os.path.join(output_dir, f"{model_name}_responses.jsonl")
    
    response_entry = {
        'task_id': task_id,
        'completion': response
    }
    
    with open(filename, 'a') as f:
        f.write(json.dumps(response_entry) + '\n')
        
@weave.op
def generate_final_response(prompt, responses, critiques):
    combined_responses = '\n\n'.join([f"Response from {resp['model']}:\n{resp['response']}" for resp in responses])
    combined_critiques = '\n\n'.join([f"Critique from {critique['model']}:\n{critique['critique']}" for critique in critiques])

    system_prompt = (
        f"You have been provided with code responses and several critiques from different models.\n"
        f"NOTE: YOU MUST PUT ALL IMPORTS INSIDE THE FUNCTION, AS ALL CODE OUTSIDE OF SINGLE FUNCTION WILL BE IGNORED."
        f"Consider these critiques to decide if any changes are needed to improve the code.\n"
        f"Handle all of the test cases and reasonable set of edge cases. DO NOT USE VALUEERRORS IN THE CODE. THIS WILL RUIN THE EVAL. JUST HANDLE THE EDGE CASES GRACEFULLY"        
        f"WRITE THE FULL COMPLETE FUNCTION (EG WITH def ....).\n"
        f"DO NOT ADD ANY EXAMPLE USES. JUST THE FUNCTION!!!!!!!!"
        f"ALL CODE MUST BE CONTAINED WITHIN THE SINGLE FUNCTION - EVEN IMPORTS ETC. YOU CAN DEFINE MORE FUNCTIONS BUT THEY MUST BE CONTAINED INSIDE THE CORE FUNCTION DESCRIBED IN THE PROMPT"
        f"YOU CANNOT USE EXTERNAL LIBRARIES LIKE NUMPY ETC. -- USE STANDARD LIBRARIES ONLY LIKE MATH, from typing import List, Tuple, from itertools import combinations, from collections import Counter, from functools import reduce, import random, import re etc etc."
        f"NOTE: YOU MUST PUT ALL IMPORTS INSIDE THE FUNCTION, AS ALL CODE OUTSIDE OF SINGLE FUNCTION WILL BE IGNORED."
        f"BE CAUTIOUS TO AVOID ' is not defined' errors too"
        f"END CODE WITH '```'. NOTE YOU ABSOLUTELY MUST END THE CODE WITH '```'.\n\n"
        f"Task:\n{prompt}\n\n"
        f"Code Responses:\n{combined_responses}\n\n"
        f"Critiques:\n{combined_critiques}\n\n"
    )
    return generate_response_from_model(system_prompt, 'openai')

def main():
    problems = read_problems()
    models = ['llama', 'openai']
    output_dir = "model_responses"
    os.makedirs(output_dir, exist_ok=True)

    completed_tasks = [task['task_id'] for task in load_progress_from_csv()]
    
    aggregated_responses = []
    for task_id, problem in problems.items():
        response = process_problem(task_id, problem, completed_tasks)
        if response:
            aggregated_responses.append(response)

        time.sleep(2)  # Delay to avoid rate limiting

    output_path = os.path.join(output_dir, "aggregated_responses_avengers.jsonl")
    write_jsonl(output_path, aggregated_responses)
    print("All aggregated responses have been saved.")

if __name__ == "__main__":
    main()