only_run_mistral.py

#model_name =  "meta-llama/CodeLlama-7b-Instruct-hf" #"mistralai/Mistral-7B-Instruct-v0.3" # "meta-llama/CodeLlama-7b-hf" #   
model_name = "deepseek-ai/deepseek-coder-6.7b-instruct" #"Qwen/Qwen2.5-Coder-7B-Instruct"  #    
# "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" #  "Qwen/Qwen2.5-7B-Instruct" # "deepseek-ai/DeepSeek-V3" # "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo" # "meta-llama/Llama-3.1-8B-Instruct" # "meta-llama/Meta-Llama-3-8B" # ""sapienzanlp/Minerva-7B-instruct-v1.0" #"google/codegemma-7b-it" #  "aiXcoder/aixcoder-7b-base" #"bigcode/starcoder2-7b" #     "google/gemma-2-9b-it" #"Qwen/Qwen2.5-Coder-7B-Instruct"  # "deepseek-ai/deepseek-coder-6.7b-instruct"  #"deepseek-ai/deepseek-coder-6.7b-base" #"deepseek-ai/DeepSeek-R1-Distill-Llama-8B" #"Qwen/Qwen2.5-Coder-7B-Instruct"  #"Qwen/Qwen2.5-Coder-7B-Instruct" #"THUDM/codegeex2-6b" #"m-a-p/OpenCodeInterpreter-DS-6.7B" #"WisdomShell/CodeShell-7B" #"NTQAI/Nxcode-CQ-7B-orpo" #"Artigenz/Artigenz-Coder-DS-6.7B" #"deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct"  # #"bigcode/starcoder2-7b" #"google/gemma-2-9b-it" #  "google/codegemma-7b-it"  #"google/gemma-7b" 
#model_name = "bigcode/starcoder2-7b" #"microsoft/codebert-base"
#model_name = "Phind/Phind-CodeLlama-34B-v2" #  "tiiuae/falcon-40b"
#out_name = "./HPC_2_mistral_iffp_20k_4_lora/checkpoint_0_10000/"
from datasets import load_dataset
from tqdm import tqdm
import json
import re
import pandas as pd
import datetime
import sys
import gc
make_dataset = False
#dataset_file = "./datasets/codecontests/diff_py_900_unperturbed.json" 
main_fname ='CC_900_py_catmix_5rand_1fold_test_conctest_javalang31_pairs_0'  #'CC_900_py_catmix_5rand_4fold_conctest_iffp_pairs_0_fulltryrem' #'CC_900_py_catmix_5rand_1fold_test_conctest_cpplang_27_oripairs' #  "CC_900_py_catmix_5rand_4fold_pairs_0_temp_fulltry" # "CC_900_py_catmix_5rand_1fold_conctest_ifelse_1s_negs_pairs_0_fulltry" # "CC_900_py_catmix_5rand_1fold_conctest_ifelse_1s_pairs_0" # #'CC_900_py_catmix_5rand_1fold_test_conctest_parallel_all_32_oripairs' #'CC_900_py_catmix_5rand_1fold_test_conctest_all_sptask_12_pairs_0' # 'CC_1900_py_catmix_5rand_1fold_conctest_all_1s_pairs_0_fulltry' # "CC_900_py_catmix_5rand_1fold_conctest_expr_0s_4_pairs_0_fulltry" #"CC_1200_py_catmix_5rand_4fold_test_1s_v_oripairs_fulltry" # "CC_1600_py_catmix_5rand_4fold_test_all_v2_oripairs_fulltry"  #"CC_2200_oripairs" # 'CC_1900_py_catmix_5rand_1fold_conctest_all_1s_oripairs_fulltry' #'CC_boolinv_pairs_0' #"CC_900_py_catmix_exact9rand_1fold_conctest_all_oripairs_fulltry" #"CC_900_py_catmix_5rand_4fold_orixpairs_temp_fulltry" #'CC_900_py_catmix_5rand_4fold_test_varc2_pairs_0_temp_fulltry2' # 'CC_900_py_catmix_5rand_1fold_conctest_k7_all_pairs_0_temp_fulltry' # 'CC_900_py_catmix_5rand_1fold_conctest_k9_all_oripairs_temp_fulltry'# "CC_2200_py_catmix_5rand_1fold_test_1s_5_pairs_0_temp_fulltry" #    "CC_900_py_catmix_5rand_1fold_conctest_all_ind_pairs_0_fulltry" # 
#'CC_900_py_catmix_3rand_1fold_conctest_all_pairs_0_temp_fulltry' #"CC_varrename_oripairs" # "CC_900_py_catmix_5rand_1fold_conctest_ifelse_1s_negs_pairs_0_fulltry" #"EquiBench_ojv" #   "CC_900_py_catmix_5rand_1fold_conctest_boolinv_1s_oripairs_fulltry" # "CC_900_py_catmix_5rand_1fold_conctest_ifelse_all_pairs_0" #"CC_2200_py_catmix_5rand_1fold_test_1s_5_oripairs_fulltry" #   "CC_900_py_catmix_5rand_4fold_test_varc2_pairs_0_temp_fulltry2" # 'CC_900_py_catmix_5rand_4fold_val_oripairs_fulltry2' #  "CN_py800_catmix_5rand_4fold_pairs_0_fulltry" #"CN_py800_catmix_5rand_4fold_full_oripairs_matched_fulltry" # "CN_py800_catmix_5rand_4fold_2000_pairs_0" #
dataset_file = f"./datasets/codecontests/{main_fname}.json"
#dataset_file = "../dataset_lfs/CC_900_py_catmix_5rand_3_pairss.json"
#output_file = "./results/31-07-mistral_iffp_900_lora4_correct_0_10000.json"
#dataset_file = "../dataset_lfs/cor_pairs_data_large_900_ind.json"
#output_file = "./results/20-08-mistral_diff_all_900_checkpoint_3_2000_use_3loop2_promptmix.json"
#output_file = f"./results/2025/18-07-mistral_{main_fname}_11071rand_0s_chat_0_10000.json" #0704_0s_both_1_4000_strictinfer.json" 
output_file = f"./results/2025/18-07-dsc_{main_fname}_base_0s_chat.json" #0704_0s_both_1_4000_strictinfer.json" 
if 'mistral' in output_file:
    model_name = "mistralai/Mistral-7B-Instruct-v0.3"
load_path = None
#load_path = "../HPC_2_codellama_iffp_20k_4_lora_900/checkpoint_1_0"
#load_path = "../HPC_2_mistral_iffp_20k_2_lora_900_loop2/checkpoint_0_14000/"
#load_path = "../HPC_13_8_mistral_var_rename_one_20k_lora_900/checkpoint_0_18000/"
#load_path = "../HPC_31_7_mistral_iffp2_20k_lora_900/checkpoint_1_16000/"
#load_path = "../HPC_09_9_mistral_var_rename_one_loop2_mix_20k_lora_900/checkpoint_1_6000/"
#load_path = "../HPC_04_9_mistral_var_rename_one_loop2_iffp2_ind_mix_20k_lora_900/checkpoint_0_60000"
#load_path = "../HPC_10_9_mistral_unperturbed_900_5e7/checkpoint_4_0/"
#load_path = "../HPC_25_9_mistral_CC_catiffp_20k_lora_900/checkpoint_0_14000"
#load_path = "../HPC_22_12_deepseekcoder_CN_py800_catmix_5rand_900_1e6/checkpoint_0_2000"
#load_path = '../HPC_20_12_qwen_CN_py800_catmix_5rand_900_1e6/checkpoint_2_2000'
#load_path = '../HPC_20_12_mistral_CN_py800_catmix_5rand_900_1e6/checkpoint_4_4000'
#load_path = '../HPC_21_01_25_mistral_CN_py800_catmix_5rand_oripairs_900_1e6_2/checkpoint_0_2000'
#load_path = '../HPC_21_01_25_qwen_CN_py800_catmix_5rand_both_900_1e6/checkpoint_0_18000'
#load_path = '../HPC_21_01_25_mistral_CN_py800_catmix_5rand_both_900_1e6_4tg/checkpoint_0_22000/'
#load_path = '../HPC_20_12_codellama_CN_py800_catmix_5rand_900_1e6/checkpoint_0_6000/'
#load_path = '../HPC_22_12_codellama_CN_py800_catmix_5rand_both_900_1e6/checkpoint_0_2000/'
#load_path = '../HPC_28_01_25_qwen_CN_py800_catmix_5rand_900_1e6_2tg/checkpoint_1_8000/'
#load_path = '../HPC_04_02_25_mistral_CN_py800_catmix_5rand_900_1e6_oribalanced/checkpoint_1_0'
#load_path = '../HPC_19_02_25_qwen_CN_py800_catmix_5rand_perchat_2/best_model/'
#load_path = '../HPC_18_03_25_mistral_CC_900_py_catmix_5rand_per/checkpoint_1_0' #0_2000/'
#load_path = '../HPC_18_03_25_qwen_CC_900_py_catmix_5rand_per/checkpoint_0_6000/'
#load_path = '../HPC_18_03_25_mistral_CC_900_py_catmix_5rand_ori/checkpoint_1_0/'
#load_path = '../HPC_18_03_25_qwen_CC_900_py_catmix_5rand_ori/checkpoint_1_6000/'
#load_path = '../HPC_18_03_25_dsc_instruct_CC_900_py_catmix_5rand_per/checkpoint_0_4000/'
#load_path = '../HPC_25_03_25_qwen_CC_900_py_catmix_5rand_per_eqm/checkpoint_1_0/'
#load_path = '../HPC_18_03_25_dsc_instruct_CC_900_py_catmix_5rand_per/checkpoint_1_2000/'
#load_path = '../HPC_25_03_25_mistral_CC_900_py_catmix_5rand_per/checkpoint_1_0/'
#load_path = '../HPC_31_03_25_qwen_CC_900_py_catmix_5rand_per_seed42_correctval/checkpoint_0_6000'
#load_path = '../HPC_31_03_25_qwen_CC_900_py_catmix_5rand_both_seed42/checkpoint_0_14000/' #mistral
#load_path = '../HPC_25_03_25_qwen_CC_900_py_catmix_5rand_per_seed42/checkpoint_0_4000/'
#load_path = '../HPC_07_04_25_mistral_CC_900_py_catmix_5rand_both_seed42/checkpoint_1_4000/'
#load_path = '../HPC_13_04_25_dsc_CC_900_py_catmix_5rand_both_seed42/checkpoint_0_14000/'
#load_path = '../HPC_17_05_25_dsc_CC_900_py_catmix_5rand_per_seed42_strict/checkpoint_1_8000/'
#load_path = '../HPC_08_07_25_qwen_CC_900_py_catmix_5rand_both_seed42_strict/checkpoint_3_6000/'
#load_path = '../HPC_11_07_25_qwen_CC_900_py_catmix_5rand_iffpboth_seed42_strict/checkpoint_1_6000/'
#load_path = '../HPC_11_07_25_dsc_CC_900_py_catmix_5rand_iffpboth_seed42_strict/checkpoint_0_10000/' #1_2000/' #
#load_path  = '../HPC_04_06_25_qwen_CC_900_py_catmix_5rand_per_seed42_strict/checkpoint_0_6000/'
#load_path = '../HPC_11_07_25_mistral_CC_900_py_catmix_5rand_iffpboth_seed42_strict/checkpoint_0_10000/'
print(load_path, dataset_file, "\n" , output_file )
start_idx = 0
#start_idx = int(sys.argv[1])
prompt_file_path =  './prompt_only_instruct_strict_dsc.txt' #'./prompt_strict_cot_one.txt' #'prompt_cot_strict_reason.txt' #'./prompt_cot_with_per_info_0s.txt' # './prompt_cot_strict_symb.txt' #'./prompt_only_with_per_info_0s.txt' #  './prompt_only_instruct_strict.txt' # './prompt_equibench.txt' #  './prompt_cot_strict_onenosymb.txt' #  './prompt_only_instruct_sym.txt' # './prompt_instruct_only_distill.txt' #'./prompt_cot_strict.txt' #  prompt_pure_only_instruct_cl.txt' #'../prompt_file.txt' #'./prompt_fewshot_chat.txt'  #'./prompt_0s.txt' #'../promp_file.txt' #'./prompt_cot.txt' #'./prompt_pure_only_instruct_cl.txt' #'./prompt_fewshot.txt'  # './prompt_pure_only_instruct_cl.txt' #'./prompt_cot_fewshot.txt'  # './prompt_pure_only_instruct.txt' #'../prompt_mix.txt' #'../prompt_only_instruct.txt' #'../prompt_mix.txt' #"../prompt_var_rename.txt" #
log_data = {
        "dataset_file": dataset_file,
        "output_file": output_file,
        "prompt_file_path": prompt_file_path,
        "model_name": model_name,
        "load_path": load_path,
        "start_idx" : start_idx,
    }


with open('onlyrun_logs.txt', 'a') as log_file:
    log_file.write(f"{datetime.datetime.now()}: running only_run_mistral.py\n")
    for key, value in log_data.items():
        log_file.write(f"{key}={value}\n")
MAX_TOK =50 #   1024 # 512 # 
BATCH_SIZE = 1
# import json
# filepath = "/kaggle/input/code-sim-try1/mutated_graph_all_lang_eq.json"
# examples = []
# with open(filepath, 'r') as file:
    # for l in file:
        # examples.append(json.loads(l))

# len(examples)
# examples[0]

#instruct_tune_dataset = load_dataset("mosaicml/instruct-v3")


def create_prompt(sample):
    bos_token = "<s>"
    original_system_message = "Below is an instruction that describes a task. Write a response that appropriately completes the request."
    system_message = "Use the provided input to create an instruction that could have been used to generate the response with an LLM."
    response = sample["prompt"].replace(original_system_message, "").replace("\n\n### Instruction\n", "").replace("\n### Response\n", "").strip()
    input = sample["response"]
    eos_token = "</s>"

    full_prompt = ""
    full_prompt += bos_token
    full_prompt += "### Instruction:"
    full_prompt += "\n" + system_message
    full_prompt += "\n\n### Input:"
    full_prompt += "\n" + input
    full_prompt += "\n\n### Response:"
    full_prompt += "\n" + response
    full_prompt += eos_token

    return full_prompt
    
    
   
   
#print(create_prompt(instruct_tune_dataset["train"][0]))



import os
import torch
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    HfArgumentParser,
    TrainingArguments,
    pipeline,
    logging,
    GemmaTokenizer,
)
from peft import LoraConfig, PeftModel
from trl import SFTTrainer

nf4_config = BitsAndBytesConfig(
   load_in_4bit=True,
   bnb_4bit_quant_type="nf4",
   bnb_4bit_use_double_quant=True,
   bnb_4bit_compute_dtype=torch.bfloat16
)

# model = AutoModelForCausalLM.from_pretrained(
    # model_name,
    # device_map='auto',
    # quantization_config=nf4_config,
    # use_cache=False
# )
if load_path is None:
    model = AutoModelForCausalLM.from_pretrained(
        #load_path,
        model_name,
        use_cache=True,
        #cache_dir = "../aib222688.scratch/HF/",
        #attn_implementation="flash_attention_2",
        torch_dtype=torch.bfloat16,
        trust_remote_code = True,
        #device_map='auto',
        #quantization_config=nf4_config,
        #use_cache=False
    )
else:

    model = AutoModelForCausalLM.from_pretrained(
        load_path,
        #model_name,
        use_cache=True,
        #cache_dir = "../aib222688.scratch/HF/",
        attn_implementation="flash_attention_2",
        torch_dtype=torch.bfloat16,
        #device_map='auto',
        #quantization_config=nf4_config,
        #use_cache=False
    )
if "gemma" in model_name:
    tokenizer = GemmaTokenizer.from_pretrained(model_name)
else:
    tokenizer = AutoTokenizer.from_pretrained(model_name)
#tokenizer = AutoTokenizer.from_pretrained(load_path)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"
# if not (load_path is None):
    # model = AutoModelForCausalLM.from_pretrained(
        # load_path
    # )

print(f"{load_path} loaded!")
# peft_config = LoraConfig(
    # lora_alpha=16,
    # lora_dropout=0.1,
    # r=64,
    # bias="none",
    # task_type="CAUSAL_LM"
# )

# import peft
# model = peft.prepare_model_for_kbit_training(model)
# model = peft.get_peft_model(model, peft_config)

# args = TrainingArguments(
  # output_dir = out_name,
  # #num_train_epochs=5,
  # max_steps = 100, 
  # per_device_train_batch_size = 4,
  # warmup_steps = 0.03,
  # logging_steps=10,
  # save_strategy="epoch",
  # #evaluation_strategy="epoch",
  # evaluation_strategy="steps",
  # eval_steps=20,
  # learning_rate=2e-4,
  # bf16=True,
  # lr_scheduler_type='constant',
# )

# max_seq_length = 2048

# trainer = SFTTrainer(
  # model=model,
  # peft_config=peft_config,
  # max_seq_length=max_seq_length,
  # tokenizer=tokenizer,
  # packing=True,
  # formatting_func=create_prompt,
  # args=args,
  # train_dataset=instruct_tune_dataset["train"],
  # eval_dataset=instruct_tune_dataset["test"]
# )


# import time
# start = time.time()
# try:
    # pass
    # #trainer.train()
# except Exception as e:
    # pass
# print(time.time()- start)


    
    
# Define the file path

#oi = "Given two programs, check if they both solve the same problem or not i.e. for any given input, would they both give the same output or not. Give the answer as either 'Yes' or 'No'. Explanation is not required."
oi = """You are a strict judge. Given two programs, you must check if they solve the same problem — that is, for any input, they must produce identical outputs.
Your task is to answer only with one word: Yes or No.
No explanation, no extra words, no formatting.
"""

prompt_ex = [("""x = input()
y = input()
z = ''
for i in range(len(x)):
    if x[i] < y[i]:
        print("-1")
        exit(0)
    else:
        z = z + y[i]
print(z)""",
"""from sys import exit
x=str(input())
y=str(input())
for i in range(len(x)):
    if x[i]<y[i]:
        print(-1)
        exit(0)
print(y)""", "Yes"),
("""x = input()
y = input()
z = ''
i = 0
while i <= len(x):
    if x[i] < y[i]:
        print('-1')
        exit(0)
    else:
        z = z + y[i]
    i += 1
print(z)""",
"""from sys import exit
x=str(input())
y=str(input())
for i in range(len(x)):
    if x[i]<y[i]:
        print(-1)
        exit(0)
print(y)""", "No"),
( """x = input()
y = input()
z = ''
for i in range(len(x)):
    if x[i] < y[i]:
        print("-1")
        exit(0)
    else:
        z = z + y[i]
print(z)""",
"""from sys import exit
ejgvpf = str(input())
y = str(input())
for i in range(len(ejgvp)):
    if ejgvp[i] < y[i]:
        print(-1)
        exit(0)
print(y)""", "No"),
("""x = input()
y = input()
z = ''
for i in range(len(x)):
    if x[i] < y[i]:
        print("-1")
        exit(0)
    else:
        z = z + y[i]
print(z)""",
"""from sys import exit
x = str(input())
y = str(input())
i = 0
while i < len(x):
    if x[i]<y[i]:
        print(-1)
        exit(0)
    i += 1
print(y)""", "Yes"),
("""x = input()
y = input()
z = ''
for i in range(len(x)):
    if x[i] < y[i]:
        print("-1")
        exit(0)
    else:
        z = z + y[i]
print(z)""",
""""x = input()
y = input()
z = ''
for i in range(len(x)):
    if x[i] >= y[i]:
        print("-1")
        exit(0)
    else:
        z = z + y[i]
print(z)""", "No" ),

( """x = input()
y = input()
z = ''
for i in range(len(x)):
    if x[i] < y[i]:
        print("-1")
        exit(0)
    else:
        z = z + y[i]
print(z)""",
"""misoc = input()
y = input()
z = ''
for i in range(len(misoc)):
    if misoc[i] < y[i]:
        print("-1")
        exit(0)
    else:
        z = z + y[i]
print(z)""", "Yes"),
("""x = input()
y = input()
z = ''
for i in range(len(x)):
    if x[i] < y[i]:
        print("-1")
        exit(0)
    else:
        z = z + y[i]
print(z)""",
"""misoc = input()
y = input()
z = ''
for i in range(len(misoc)):
    if miso[i] < y[i]:
        print("-1")
        exit(0)
    else:
        z = z + y[i]
print(z)""", "No" ),
("""x = input()
y = input()
z = ''
for i in range(len(x)):
    if x[i] < y[i]:
        print("-1")
        exit(0)
    else:
        z = z + y[i]
print(z)""",
"""def myfun(s1, s2):
    m, n = len(s1), len(s2)
    dp = [[0] * (n + 1) for _ in range(m + 1)]
    for i in range(1, m + 1):
        for j in range(1, n + 1):
            if s1[i - 1] == s2[j - 1]:
                dp[i][j] = dp[i - 1][j - 1] + 1
            else:
                dp[i][j] = max(dp[i - 1][j], dp[i][j - 1])
    ans = []
    i, j = m, n
    while i > 0 and j > 0:
        if s1[i - 1] == s2[j - 1]:
            ans.append(s1[i - 1])
            i -= 1
            j -= 1
        elif dp[i - 1][j] >= dp[i][j - 1]:
            i -= 1
        else:
            j -= 1

    return ''.join(reversed(ans))""", "No"),
("""x = input()
y = input()
z = ''
for i in range(len(x)):
    if x[i] < y[i]:
        print("-1")
        exit(0)
    else:
        z = z + y[i]
print(z)""",
"""x = input()\ny = input()\n\nalf = 'abcdefghijklmnopqrstuvwxyz'\n\nz = ''\n\nfor i in range(len(x)):\n\tif y[i] == x[i]:\n\t\tz += x[i]\n\telif alf.find(x[i]) < alf.find(y[i]):\n\t\tprint('-1')\n\t\tbreak\n\telif alf.find(x[i]) > alf.find(y[i]):\n\t\tz += y[i]\n\nelse:\n\tprint(z)""", "Yes"),
("""def check(s1,s2):
    l = len(s1)
    for i in range(l):
        if s1[i] >= s2[i]:
            continue
        else:
            return False
    return True


x = input()
y = input()

if check(x,y):
    print(y)
else:
    print(-1)""", 
"""def check(s1,s2):
    l = len(s1)
    for i in range(l):
        if s1[i] < s2[i]:
            return False
        else:
            continue
    return True


x = input()
y = input()

if check(x,y):
    print(y)
else:
    print(-1)"""    , "Yes")]
    
def format_few_shot_prompt( nprog1, nprog2, model_name='codellama', overarching_instruction=oi, examples=prompt_ex):
    """
    Formats a few-shot prompt based on the specified model's instruction format.

    Parameters:
    - model_name (str): The name of the model (e.g., 'mistral-7b-instruct', 'qwen-7b-instruct', 'codellama-7b-instruct').
    - overarching_instruction (str): The high-level instruction to apply to all examples.
    - examples (list of tuples): A list where each tuple contains (input_text, output_text).
    - new_input (str): The new input text for which the model should generate a response.

    Returns:
    - str: The formatted prompt.
    """
    if 'mistral' in model_name.lower():
        prompt = ""
        for prog1, prog2, output_text in examples:
            prompt += f"<s>[INST] {overarching_instruction}\nProgram 1: {prog1}\nProgram 2: {prog2} [/INST] {output_text} </s>\n"
        prompt += f"<s>[INST] {overarching_instruction}\nProgram 1: {nprog1}\nProgram 2: {nprog2} [/INST] "

    elif 'qwen' in model_name.lower():
        prompt = f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
        for prog1, prog2, output_text in examples:
            prompt += f"<|im_start|>user\n{overarching_instruction}\nProgram 1: {prog1}\nProgram 2: {prog2}<|im_end|>\n"
            prompt += f"<|im_start|>assistant\n{output_text}<|im_end|>\n"
        prompt += f"<|im_start|>user\n{overarching_instruction}\nProgram 1: {nprog1}\nProgram 2: {nprog2}<|im_end|>\n"
        prompt += f"<|im_start|>assistant\n"

    elif 'codellama' in model_name.lower():
        prompt = f"### Instruction:\n{overarching_instruction}\n\n"
        for prog1, prog2, output_text in examples:
            prompt += f"### Instruction:\n{overarching_instruction}\nProgram 1: {prog1}\nProgram 2: {prog2}\n\n### Response:\n{output_text}\n\n"
        prompt += f"### Instruction:\n{overarching_instruction}\nProgram 1: {nprog1}\nProgram 2: {nprog2}\n\n### Response:\n"

    else:
        raise ValueError("Model name not recognized or unsupported.")

    return prompt




# formatted_prompt = format_few_shot_prompt("""p1""", 'p2', 'codellama')
# print(formatted_prompt)

def generate_batch(prompts, model):
    try:
        batch_messages = []
        if '_0s_' in output_file and 'chat' in output_file:
            # CHAT TEMPLATE path
            for prompt in prompts:
                messages = [
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": prompt},
                ]
                batch_messages.append(messages)

            tokenized_chats = tokenizer.apply_chat_template(
                batch_messages,
                tokenize=True,
                add_generation_prompt=True,
                return_tensors="pt",
                padding=True,
                truncation=True
            )

            # --- ADD ATTENTION MASK ---
            model_inputs = {
                "input_ids": tokenized_chats.to('cuda'),
                "attention_mask": (tokenized_chats != tokenizer.pad_token_id).to('cuda')  # <== create attention_mask
            }
            
            generated_ids = model.generate(
                **model_inputs,
                max_new_tokens=MAX_TOK,
                do_sample=False,
                pad_token_id=tokenizer.eos_token_id
            )

            decoded_outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

            clean_outputs = []
            for idx, output in enumerate(decoded_outputs):
                input_text = tokenizer.decode(tokenized_chats[idx], skip_special_tokens=True)
                clean_outputs.append(output.replace(input_text, ""))
            
            return clean_outputs

        else:
            # NORMAL ENCODING path
            encoded_inputs = tokenizer(
                prompts,
                return_tensors="pt",
                padding=True,
                truncation=True,
                add_special_tokens=True
            )

            model_inputs = {k: v.to('cuda') for k, v in encoded_inputs.items()}

            generated_ids = model.generate(
                input_ids=model_inputs["input_ids"],
                attention_mask=model_inputs["attention_mask"],  # <== explicitly pass
                max_new_tokens=MAX_TOK,
                do_sample=False,
                pad_token_id=tokenizer.eos_token_id
            )

            decoded_outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

            clean_outputs = []
            for prompt, output in zip(prompts, decoded_outputs):
                clean_outputs.append(output.replace(prompt, ""))
            
            return clean_outputs
    except torch.cuda.OutOfMemoryError:
            print("CUDA ran out of memory! Clearing cache and retrying...")
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
            gc.collect()
            return None
            
    except torch.OutOfMemoryError:
            print("CUDA ran out of memory! Clearing cache and retrying...")
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
            gc.collect()
            with open(output_file, 'a+') as outfile:
                for response in responses:
                    outfile.write(f"{json.dumps(response)}\n")
            responses = []
            return None


if 'chat' in output_file:

    def generate_response(prompt, model):
        try:
            
            # if "Nxcode" in model_name:
                # del model_inputs['token_type_ids']
            messages = [ {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt}, ]
            if '_0s_' in output_file: 
                tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
                model_inputs = tokenized_chat.to('cuda')
                # model.to("cpu")
                # model_inputs = model_inputs.to("cpu")
                if len(model_inputs)>2000:
                    print(len(model_inputs))
                generated_ids = model.generate(model_inputs, max_new_tokens=MAX_TOK, do_sample=False, pad_token_id=tokenizer.eos_token_id)
                decoded_output = tokenizer.batch_decode(generated_ids)
                return decoded_output[0].replace(tokenizer.decode(tokenized_chat[0]), "")

            else:
                #print(prompt)
                encoded_input = tokenizer(prompt,  return_tensors="pt", add_special_tokens=True)
                model_inputs = encoded_input.to('cuda')
                
            #print(tokenized_chat, tokenizer.apply_chat_template(messages, tokenize=False,  add_generation_prompt=True,))
            generated_ids = model.generate(**model_inputs, max_new_tokens=MAX_TOK, do_sample=False, pad_token_id=tokenizer.eos_token_id)
            decoded_output = tokenizer.batch_decode(generated_ids)
            return decoded_output[0].replace(prompt, "")
            
            # torch.cuda.empty_cache()
            # torch.cuda.ipc_collect()
            
            #print(decoded_output[0].replace(tokenizer.decode(tokenized_chat[0]), "")) 
            
            
            
        except torch.cuda.OutOfMemoryError:
            print("CUDA ran out of memory! Clearing cache and retrying...")
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
            gc.collect()
            return None
            
        except torch.OutOfMemoryError:
            print("CUDA ran out of memory! Clearing cache and retrying...")
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
            gc.collect()
            with open(output_file, 'a+') as outfile:
                for response in responses:
                    outfile.write(f"{json.dumps(response)}\n")
            responses = []
            return None




    
  
else:
    
    #prompt_file_path = './prompt_0s.txt' # '../prompt_file.txt' #'./prompt_0s.txt' #
    def generate_response(prompt, model):
        try:
            #prompt = prompt + "< [/INST]</s>>"
            encoded_input = tokenizer(prompt,  return_tensors="pt", add_special_tokens=True)
            model_inputs = encoded_input.to('cuda')
            if "xcode" in model_name:
                del model_inputs['token_type_ids']
            generated_ids = model.generate(**model_inputs, max_new_tokens=MAX_TOK, do_sample=False, pad_token_id=tokenizer.eos_token_id)

            decoded_output = tokenizer.batch_decode(generated_ids)
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
            #print(decoded_output[0].replace(prompt, ""))
            return decoded_output[0].replace(prompt, "")
            
        except torch.cuda.OutOfMemoryError:
            print("CUDA ran out of memory! Clearing cache and retrying...")
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
            gc.collect()
            return None
                
        except torch.OutOfMemoryError:
            print("CUDA ran out of memory! Clearing cache and retrying...")
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
            gc.collect()
            with open(output_file, 'a+') as outfile:
                for response in responses:
                    outfile.write(f"{json.dumps(response)}\n")
            responses = []
            return None





device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#model = nn.DataParallel(model, device_ids=[0, 1])
model.to(device)
model.eval()
print(model)
#file_path = '../prompt_file.txt'


# Open the file in read mode and store the content as a string
with open(prompt_file_path, 'r') as file:
    file_content = file.read()


allpairs = []


# with open(dataset_file, 'r') as file:
    # for l in file:
        # allpairs.append(json.loads(l))
        
with open(dataset_file, 'r') as file:
    for l in file:
        fobj = json.loads(l)
        if not (fobj in allpairs):
            allpairs.append(fobj)

input_ids = tokenizer("Hello world", return_tensors="pt").input_ids.to("cuda")
output = model.generate(input_ids)
print('o/p: ', output)

if BATCH_SIZE==1:
    try:
        responses = []
        allpairs = allpairs[start_idx:]
        first = 5
        for pair in tqdm(allpairs):
            
            
            
            prompt = file_content + pair['prog1']['scode'] + "\nProgram 2:"  #   "\n[Program 2]:" #
            prompt += pair['prog2']['scode'] + "\n"
            true_label = 'YES'
            false_label = 'NO'
            #prompt += f"Please only output the answer of whether the two programs are equivalent or not. You should only output {true_label} or {false_label}."
            if '!!!!' in output_file:
                prompt = format_few_shot_prompt(pair['prog1']['scode'], pair['prog1']['scode'], model_name)

            if 'chat' not in output_file:
                pass
            #prompt += "###Response:"
            result=generate_response(prompt, model)
            if first> 0:
                print(prompt, result)
                first -= 1
            end_ind = 0
            if result:
                end_ind = result.find("###")
            else:   
                result = ""
                res = ""
            res = result
            if end_ind>0:
                res = result[:end_ind]
            ans = pair['prog1']['probid'] == pair['prog2']['probid']
            correct = False
            if ans:
                if "Yes" in res or "yes" in res:
                    correct = True
            if ans == False:
                if "No" in res or "no" in res:
                    correct = True
           
            responses.append({
                'progs' : pair,
                'response': res,
                'extracted_answer': correct,
                'full' : result,
            })
            
        
        with open(output_file, 'a+') as outfile:
            for response in responses:
                outfile.write(f"{json.dumps(response)}\n")
            responses = []
            
            
    
    except KeyboardInterrupt:
        with open(output_file, 'a+') as outfile:
            for response in responses:
                outfile.write(f"{json.dumps(response)}\n")
            print(len(responses))
            responses = []
        

else:
    try:
        responses = []
        allpairs = allpairs[start_idx:]
        first = 2
        for i in tqdm(range(0, len(allpairs), BATCH_SIZE)):
            batch = allpairs[i:i+BATCH_SIZE] 
            prompts = []
            for pair in batch:
            #if pair['category'].startswith("0_1_"):
                prompt = file_content + pair['prog1']['scode'] + "\nProgram 2:" 
                prompt += pair['prog2']['scode'] + "\n"
                if '_0s_' not in output_file:
                    prompt = format_few_shot_prompt(pair['prog1']['scode'], pair['prog1']['scode'], model_name)

                if 'chat' not in output_file:
                    pass
                #prompt += "###Response:"
                prompts.append(prompt)
                
            #result=generate_response(prompt, model)
            results = generate_batch(prompts, model)
            if first> 0:
                print(prompts, results)
                first -= 1
            end_ind = 0
            if results is None:
                i -= 1
                continue
            for rid, result in enumerate(results):
                if result:
                    end_ind = result.find("###")
                else:   
                    result = ""
                    res = ""
                res = result
                if end_ind>0:
                    res = result[:end_ind]
                ans = pair['prog1']['probid'] == pair['prog2']['probid']
                correct = False
                if ans:
                    if "Yes" in res or "yes" in res:
                        correct = True
                if ans == False:
                    if "No" in res or "no" in res:
                        correct = True
               
                responses.append({
                    'progs' : batch[rid],
                    'response': res,
                    'extracted_answer': correct,
                    'full' : result,
                })
            
        
        with open(output_file, 'a+') as outfile:
            for response in responses:
                outfile.write(f"{json.dumps(response)}\n")
            responses = []
        
   
    

    except KeyboardInterrupt:
        with open(output_file, 'a+') as outfile:
            for response in responses:
                outfile.write(f"{json.dumps(response)}\n")
            print(len(responses))
            responses = []
        

# except Exception as e:
    # print(f"TPR: {TP/(TP+FN)}\nFPR: {TN/(TN+FP)} \nAcc: {(TN+TP)/(TP+FP+TN+FN)}  ")
with open(output_file, 'a+') as outfile:
    for response in responses:
        outfile.write(f"{json.dumps(response)}\n")
    print(len(responses))
    responses = []
print(len(responses))


# model_name =  "meta-llama/CodeLlama-7b-Instruct-hf" #"meta-llama/CodeLlama-7b-hf" #"mistralai/Mistral-7B-Instruct-v0.3" #    
# model_name = "Qwen/Qwen2.5-Coder-7B-Instruct" #"deepseek-ai/deepseek-coder-6.7b-instruct"  #"deepseek-ai/deepseek-coder-6.7b-base" #"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" #"deepseek-ai/DeepSeek-R1-Distill-Llama-8B" #"Qwen/Qwen2.5-Coder-7B-Instruct"  #"Qwen/Qwen2.5-Coder-7B-Instruct" #"google/gemma-2-9b-it" #"THUDM/codegeex2-6b" #"m-a-p/OpenCodeInterpreter-DS-6.7B" #"WisdomShell/CodeShell-7B" #"google/codegemma-7b-it" #"NTQAI/Nxcode-CQ-7B-orpo" #"Artigenz/Artigenz-Coder-DS-6.7B" #"meta-llama/Meta-Llama-3-8B" #"deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct"  # #"bigcode/starcoder2-7b" #"google/gemma-2-9b-it" #  "google/codegemma-7b-it"  #"google/gemma-7b" "bigcode/starcoder2-7b" #
# #model_name = "bigcode/starcoder2-7b" #"microsoft/codebert-base"
# #model_name = "Phind/Phind-CodeLlama-34B-v2" #  "tiiuae/falcon-40b"
# #out_name = "./HPC_2_mistral_iffp_20k_4_lora/checkpoint_0_10000/"
# from datasets import load_dataset
# from tqdm import tqdm
# import json
# import re
# import pandas as pd
# import datetime
# import sys
# import gc
# make_dataset = False
# #dataset_file = "./datasets/codecontests/diff_py_900_unperturbed.json" 
# main_fname = "CC_900_py_catmix_5rand_4fold_pairs_0_temp_fulltry" #"CN_py800_catmix_5rand_4fold_pairs_0_fulltry" #"CN_py800_catmix_5rand_4fold_full_oripairs_matched_fulltry" #  "CC_900_py_catmix_5rand_4fold_pairs_0_temp_fulltry" #"CN_py800_catmix_5rand_4fold_2000_pairs_0" #
# dataset_file = f"./datasets/codecontests/{main_fname}.json"
# #dataset_file = "../dataset_lfs/CC_900_py_catmix_5rand_3_pairss.json"
# #output_file = "./results/31-07-mistral_iffp_900_lora4_correct_0_10000.json"
# #dataset_file = "../dataset_lfs/cor_pairs_data_large_900_ind.json"
# #output_file = "./results/20-08-mistral_diff_all_900_checkpoint_3_2000_use_3loop2_promptmix.json"
# output_file = f"./results/2025/25-03-mistral_{main_fname}_chat_0s_base2.json" 
# if 'mistral' in output_file:
    # model_name = "mistralai/Mistral-7B-Instruct-v0.3"
# load_path = None
# #load_path = "../HPC_2_codellama_iffp_20k_4_lora_900/checkpoint_1_0"
# #load_path = "../HPC_2_mistral_iffp_20k_2_lora_900_loop2/checkpoint_0_14000/"
# #load_path = "../HPC_13_8_mistral_var_rename_one_20k_lora_900/checkpoint_0_18000/"
# #load_path = "../HPC_31_7_mistral_iffp2_20k_lora_900/checkpoint_1_16000/"
# #load_path = "../HPC_09_9_mistral_var_rename_one_loop2_mix_20k_lora_900/checkpoint_1_6000/"
# #load_path = "../HPC_04_9_mistral_var_rename_one_loop2_iffp2_ind_mix_20k_lora_900/checkpoint_0_60000"
# #load_path = "../HPC_10_9_mistral_unperturbed_900_5e7/checkpoint_4_0/"
# #load_path = "../HPC_25_9_mistral_CC_catiffp_20k_lora_900/checkpoint_0_14000"
# #load_path = "../HPC_22_12_deepseekcoder_CN_py800_catmix_5rand_900_1e6/checkpoint_0_2000"
# #load_path = '../HPC_20_12_qwen_CN_py800_catmix_5rand_900_1e6/checkpoint_2_2000'
# #load_path = '../HPC_20_12_mistral_CN_py800_catmix_5rand_900_1e6/checkpoint_4_4000'
# #load_path = '../HPC_21_01_25_mistral_CN_py800_catmix_5rand_oripairs_900_1e6_2/checkpoint_0_2000'
# #load_path = '../HPC_21_01_25_qwen_CN_py800_catmix_5rand_both_900_1e6/checkpoint_0_18000'
# #load_path = '../HPC_21_01_25_mistral_CN_py800_catmix_5rand_both_900_1e6_4tg/checkpoint_0_22000/'
# #load_path = '../HPC_20_12_codellama_CN_py800_catmix_5rand_900_1e6/checkpoint_0_6000/'
# #load_path = '../HPC_22_12_codellama_CN_py800_catmix_5rand_both_900_1e6/checkpoint_0_2000/'
# #load_path = '../HPC_28_01_25_qwen_CN_py800_catmix_5rand_900_1e6_2tg/checkpoint_1_8000/'
# #load_path = '../HPC_04_02_25_mistral_CN_py800_catmix_5rand_900_1e6_oribalanced/checkpoint_1_0'
# #load_path = '../HPC_19_02_25_qwen_CN_py800_catmix_5rand_perchat_2/best_model/'
# #load_path = '../HPC_18_03_25_mistral_CC_900_py_catmix_5rand_per/checkpoint_1_0' #0_2000/'
# #load_path = '../HPC_18_03_25_qwen_CC_900_py_catmix_5rand_per/checkpoint_0_6000/'
# #load_path = '../HPC_18_03_25_mistral_CC_900_py_catmix_5rand_ori/checkpoint_1_0/'
# #load_path = '../HPC_18_03_25_qwen_CC_900_py_catmix_5rand_ori/checkpoint_1_6000/'
# #load_path = '../HPC_18_03_25_dsc_instruct_CC_900_py_catmix_5rand_per/checkpoint_0_4000/'
# #load_path = '../HPC_25_03_25_qwen_CC_900_py_catmix_5rand_per_eqm/checkpoint_1_0/'
# #load_path = '../HPC_18_03_25_dsc_instruct_CC_900_py_catmix_5rand_per/checkpoint_1_2000/'
# #load_path = '../HPC_25_03_25_mistral_CC_900_py_catmix_5rand_per/checkpoint_1_0/'
# print(load_path, dataset_file, "\n" , output_file )
# start_idx = 0
# #start_idx = int(sys.argv[1])
# prompt_file_path ='./prompt_pure_only_instruct_cl.txt' #'../prompt_file.txt' #'./prompt_fewshot_chat.txt'  #'./prompt_0s.txt' #'../promp_file.txt' #'./prompt_cot.txt' #'./prompt_pure_only_instruct_cl.txt' #'./prompt_fewshot.txt'  # './prompt_pure_only_instruct_cl.txt' #'./prompt_cot_fewshot.txt'  # './prompt_pure_only_instruct.txt' #'../prompt_mix.txt' #'../prompt_only_instruct.txt' #'../prompt_mix.txt' #"../prompt_var_rename.txt" #
# log_data = {
        # "dataset_file": dataset_file,
        # "output_file": output_file,
        # "prompt_file_path": prompt_file_path,
        # "model_name": model_name,
        # "load_path": load_path,
        # "start_idx" : start_idx,
    # }


# with open('onlyrun_logs.txt', 'a') as log_file:
    # log_file.write(f"{datetime.datetime.now()}: running only_run_mistral.py\n")
    # for key, value in log_data.items():
        # log_file.write(f"{key}={value}\n")
# MAX_TOK = 50 #12
# # import json
# # filepath = "/kaggle/input/code-sim-try1/mutated_graph_all_lang_eq.json"
# # examples = []
# # with open(filepath, 'r') as file:
    # # for l in file:
        # # examples.append(json.loads(l))

# # len(examples)
# # examples[0]

# #instruct_tune_dataset = load_dataset("mosaicml/instruct-v3")


# def create_prompt(sample):
    # bos_token = "<s>"
    # original_system_message = "Below is an instruction that describes a task. Write a response that appropriately completes the request."
    # system_message = "Use the provided input to create an instruction that could have been used to generate the response with an LLM."
    # response = sample["prompt"].replace(original_system_message, "").replace("\n\n### Instruction\n", "").replace("\n### Response\n", "").strip()
    # input = sample["response"]
    # eos_token = "</s>"

    # full_prompt = ""
    # full_prompt += bos_token
    # full_prompt += "### Instruction:"
    # full_prompt += "\n" + system_message
    # full_prompt += "\n\n### Input:"
    # full_prompt += "\n" + input
    # full_prompt += "\n\n### Response:"
    # full_prompt += "\n" + response
    # full_prompt += eos_token

    # return full_prompt
    
    
   
   
# #print(create_prompt(instruct_tune_dataset["train"][0]))



# import os
# import torch
# from datasets import load_dataset
# from transformers import (
    # AutoModelForCausalLM,
    # AutoTokenizer,
    # BitsAndBytesConfig,
    # HfArgumentParser,
    # TrainingArguments,
    # pipeline,
    # logging,
    # GemmaTokenizer,
# )
# from peft import LoraConfig, PeftModel
# from trl import SFTTrainer

# nf4_config = BitsAndBytesConfig(
   # load_in_4bit=True,
   # bnb_4bit_quant_type="nf4",
   # bnb_4bit_use_double_quant=True,
   # bnb_4bit_compute_dtype=torch.bfloat16
# )

# # model = AutoModelForCausalLM.from_pretrained(
    # # model_name,
    # # device_map='auto',
    # # quantization_config=nf4_config,
    # # use_cache=False
# # )
# if load_path is None:
    # model = AutoModelForCausalLM.from_pretrained(
        # #load_path,
        # model_name,
        # use_cache=True,
        # #cache_dir = "../aib222688.scratch/HF/",
        # #attn_implementation="flash_attention_2",
        # torch_dtype=torch.float16,
        # trust_remote_code = True,
        # #device_map='auto',
        # #quantization_config=nf4_config,
        # #use_cache=False
    # )
# else:

    # model = AutoModelForCausalLM.from_pretrained(
        # load_path,
        # #model_name,
        # use_cache=True,
        # #cache_dir = "../aib222688.scratch/HF/",
        # #attn_implementation="flash_attention_2",
        # torch_dtype=torch.float16,
        # #device_map='auto',
        # #quantization_config=nf4_config,
        # #use_cache=False
    # )
# if "gemma" in model_name:
    # tokenizer = GemmaTokenizer.from_pretrained(model_name)
# else:
    # tokenizer = AutoTokenizer.from_pretrained(model_name)
# #tokenizer = AutoTokenizer.from_pretrained(load_path)
# tokenizer.pad_token = tokenizer.eos_token
# tokenizer.padding_side = "right"
# # if not (load_path is None):
    # # model = AutoModelForCausalLM.from_pretrained(
        # # load_path
    # # )

# print(f"{load_path} loaded!")
# # peft_config = LoraConfig(
    # # lora_alpha=16,
    # # lora_dropout=0.1,
    # # r=64,
    # # bias="none",
    # # task_type="CAUSAL_LM"
# # )

# # import peft
# # model = peft.prepare_model_for_kbit_training(model)
# # model = peft.get_peft_model(model, peft_config)

# # args = TrainingArguments(
  # # output_dir = out_name,
  # # #num_train_epochs=5,
  # # max_steps = 100, 
  # # per_device_train_batch_size = 4,
  # # warmup_steps = 0.03,
  # # logging_steps=10,
  # # save_strategy="epoch",
  # # #evaluation_strategy="epoch",
  # # evaluation_strategy="steps",
  # # eval_steps=20,
  # # learning_rate=2e-4,
  # # bf16=True,
  # # lr_scheduler_type='constant',
# # )

# # max_seq_length = 2048

# # trainer = SFTTrainer(
  # # model=model,
  # # peft_config=peft_config,
  # # max_seq_length=max_seq_length,
  # # tokenizer=tokenizer,
  # # packing=True,
  # # formatting_func=create_prompt,
  # # args=args,
  # # train_dataset=instruct_tune_dataset["train"],
  # # eval_dataset=instruct_tune_dataset["test"]
# # )


# # import time
# # start = time.time()
# # try:
    # # pass
    # # #trainer.train()
# # except Exception as e:
    # # pass
# # print(time.time()- start)


    
    
# # Define the file path

# oi = "Given two programs, check if they both solve the same problem or not i.e. for any given input, would they both give the same output or not. Give the answer as either 'Yes' or 'No'. Explanation is not required."

# prompt_ex = [("""x = input()
# y = input()
# z = ''
# for i in range(len(x)):
    # if x[i] < y[i]:
        # print("-1")
        # exit(0)
    # else:
        # z = z + y[i]
# print(z)""",
# """from sys import exit
# x=str(input())
# y=str(input())
# for i in range(len(x)):
    # if x[i]<y[i]:
        # print(-1)
        # exit(0)
# print(y)""", "Yes"),
# ("""x = input()
# y = input()
# z = ''
# i = 0
# while i <= len(x):
    # if x[i] < y[i]:
        # print('-1')
        # exit(0)
    # else:
        # z = z + y[i]
    # i += 1
# print(z)""",
# """from sys import exit
# x=str(input())
# y=str(input())
# for i in range(len(x)):
    # if x[i]<y[i]:
        # print(-1)
        # exit(0)
# print(y)""", "No"),
# ( """x = input()
# y = input()
# z = ''
# for i in range(len(x)):
    # if x[i] < y[i]:
        # print("-1")
        # exit(0)
    # else:
        # z = z + y[i]
# print(z)""",
# """from sys import exit
# ejgvpf = str(input())
# y = str(input())
# for i in range(len(ejgvp)):
    # if ejgvp[i] < y[i]:
        # print(-1)
        # exit(0)
# print(y)""", "No"),
# ("""x = input()
# y = input()
# z = ''
# for i in range(len(x)):
    # if x[i] < y[i]:
        # print("-1")
        # exit(0)
    # else:
        # z = z + y[i]
# print(z)""",
# """from sys import exit
# x = str(input())
# y = str(input())
# i = 0
# while i < len(x):
    # if x[i]<y[i]:
        # print(-1)
        # exit(0)
    # i += 1
# print(y)""", "Yes"),
# ("""x = input()
# y = input()
# z = ''
# for i in range(len(x)):
    # if x[i] < y[i]:
        # print("-1")
        # exit(0)
    # else:
        # z = z + y[i]
# print(z)""",
# """"x = input()
# y = input()
# z = ''
# for i in range(len(x)):
    # if x[i] >= y[i]:
        # print("-1")
        # exit(0)
    # else:
        # z = z + y[i]
# print(z)""", "No" ),
# ( """x = input()
# y = input()
# z = ''
# for i in range(len(x)):
    # if x[i] < y[i]:
        # print("-1")
        # exit(0)
    # else:
        # z = z + y[i]
# print(z)""",
# """misoc = input()
# y = input()
# z = ''
# for i in range(len(misoc)):
    # if misoc[i] < y[i]:
        # print("-1")
        # exit(0)
    # else:
        # z = z + y[i]
# print(z)""", "Yes"),
# ("""x = input()
# y = input()
# z = ''
# for i in range(len(x)):
    # if x[i] < y[i]:
        # print("-1")
        # exit(0)
    # else:
        # z = z + y[i]
# print(z)""",
# """misoc = input()
# y = input()
# z = ''
# for i in range(len(misoc)):
    # if miso[i] < y[i]:
        # print("-1")
        # exit(0)
    # else:
        # z = z + y[i]
# print(z)""", "No" ),
# ("""x = input()
# y = input()
# z = ''
# for i in range(len(x)):
    # if x[i] < y[i]:
        # print("-1")
        # exit(0)
    # else:
        # z = z + y[i]
# print(z)""",
# """def myfun(s1, s2):
    # m, n = len(s1), len(s2)
    # dp = [[0] * (n + 1) for _ in range(m + 1)]
    # for i in range(1, m + 1):
        # for j in range(1, n + 1):
            # if s1[i - 1] == s2[j - 1]:
                # dp[i][j] = dp[i - 1][j - 1] + 1
            # else:
                # dp[i][j] = max(dp[i - 1][j], dp[i][j - 1])
    # ans = []
    # i, j = m, n
    # while i > 0 and j > 0:
        # if s1[i - 1] == s2[j - 1]:
            # ans.append(s1[i - 1])
            # i -= 1
            # j -= 1
        # elif dp[i - 1][j] >= dp[i][j - 1]:
            # i -= 1
        # else:
            # j -= 1

    # return ''.join(reversed(ans))""", "No")]
    
# def format_few_shot_prompt( nprog1, nprog2, model_name='codellama', overarching_instruction=oi, examples=prompt_ex):
    # """
    # Formats a few-shot prompt based on the specified model's instruction format.

    # Parameters:
    # - model_name (str): The name of the model (e.g., 'mistral-7b-instruct', 'qwen-7b-instruct', 'codellama-7b-instruct').
    # - overarching_instruction (str): The high-level instruction to apply to all examples.
    # - examples (list of tuples): A list where each tuple contains (input_text, output_text).
    # - new_input (str): The new input text for which the model should generate a response.

    # Returns:
    # - str: The formatted prompt.
    # """
    # if 'mistral' in model_name.lower():
        # prompt = ""
        # for prog1, prog2, output_text in examples:
            # prompt += f"<s>[INST] {overarching_instruction}\nProgram 1: {prog1}\nProgram 2: {prog2} [/INST] {output_text} </s>\n"
        # prompt += f"<s>[INST] {overarching_instruction}\nProgram 1: {nprog1}\nProgram 2: {nprog2} [/INST] "

    # elif 'qwen' in model_name.lower():
        # prompt = f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
        # for prog1, prog2, output_text in examples:
            # prompt += f"<|im_start|>user\n{overarching_instruction}\nProgram 1: {prog1}\nProgram 2: {prog2}<|im_end|>\n"
            # prompt += f"<|im_start|>assistant\n{output_text}<|im_end|>\n"
        # prompt += f"<|im_start|>user\n{overarching_instruction}\nProgram 1: {nprog1}\nProgram 2: {nprog2}<|im_end|>\n"
        # prompt += f"<|im_start|>assistant\n"

    # elif 'codellama' in model_name.lower():
        # prompt = f"### Instruction:\n{overarching_instruction}\n\n"
        # for prog1, prog2, output_text in examples:
            # prompt += f"### Instruction:\n{overarching_instruction}\nProgram 1: {prog1}\nProgram 2: {prog2}\n\n### Response:\n{output_text}\n\n"
        # prompt += f"### Instruction:\n{overarching_instruction}\nProgram 1: {nprog1}\nProgram 2: {nprog2}\n\n### Response:\n"

    # else:
        # raise ValueError("Model name not recognized or unsupported.")

    # return prompt




# formatted_prompt = format_few_shot_prompt("""p1""", 'p2', 'codellama')
# print(formatted_prompt)

# if 'chat' in output_file:

    # def generate_response(prompt, model):
        # try:
            
            # # if "Nxcode" in model_name:
                # # del model_inputs['token_type_ids']
            # messages = [ {"role": "system", "content": "You are a helpful assistant."},
                # {"role": "user", "content": prompt}, ]
            # if '0s' in output_file: 
                # tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
                # model_inputs = tokenized_chat.to('cuda')
                # # model.to("cpu")
                # # model_inputs = model_inputs.to("cpu")
                # generated_ids = model.generate(model_inputs, max_new_tokens=MAX_TOK, do_sample=True, pad_token_id=tokenizer.eos_token_id)
                # decoded_output = tokenizer.batch_decode(generated_ids)
                # return decoded_output[0].replace(tokenizer.decode(tokenized_chat[0]), "")

            # else:
                # #print(prompt)
                # encoded_input = tokenizer(prompt,  return_tensors="pt", add_special_tokens=True)
                # model_inputs = encoded_input.to('cuda')
                
            # #print(tokenized_chat, tokenizer.apply_chat_template(messages, tokenize=False,  add_generation_prompt=True,))
            # generated_ids = model.generate(**model_inputs, max_new_tokens=MAX_TOK, do_sample=True, pad_token_id=tokenizer.eos_token_id)
            # decoded_output = tokenizer.batch_decode(generated_ids)
            # return decoded_output[0].replace(prompt, "")
            
            # # torch.cuda.empty_cache()
            # # torch.cuda.ipc_collect()
            
            # #print(decoded_output[0].replace(tokenizer.decode(tokenized_chat[0]), "")) 
            
            
            
        # except torch.cuda.OutOfMemoryError:
            # print("CUDA ran out of memory! Clearing cache and retrying...")
            # torch.cuda.empty_cache()
            # torch.cuda.ipc_collect()
            # gc.collect()
            # return None




    
  
# else:
    # prompt_file_path = './prompt_0s.txt' # '../prompt_file.txt' #'./prompt_0s.txt' #
    # def generate_response(prompt, model):
        # #prompt = prompt + "< [/INST]</s>>"
        # encoded_input = tokenizer(prompt,  return_tensors="pt", add_special_tokens=True)
        # model_inputs = encoded_input.to('cuda')
        # if "Nxcode" in model_name:
            # del model_inputs['token_type_ids']
        # generated_ids = model.generate(**model_inputs, max_new_tokens=MAX_TOK, do_sample=True, pad_token_id=tokenizer.eos_token_id)

        # decoded_output = tokenizer.batch_decode(generated_ids)
        # torch.cuda.empty_cache()
        # torch.cuda.ipc_collect()
        # #print(decoded_output[0].replace(prompt, ""))
        # return decoded_output[0].replace(prompt, "")



# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# #model = nn.DataParallel(model, device_ids=[0, 1])
# model.to(device)
# #file_path = '../prompt_file.txt'


# # Open the file in read mode and store the content as a string
# with open(prompt_file_path, 'r') as file:
    # file_content = file.read()


# allpairs = []


# with open(dataset_file, 'r') as file:
    # for l in file:
        # allpairs.append(json.loads(l))

# input_ids = tokenizer("Hello world", return_tensors="pt").input_ids.to("cuda")
# output = model.generate(input_ids)
# print('o/p: ', output)

        
   
    

# except KeyboardInterrupt:
    # with open(output_file, 'a+') as outfile:
        # for response in responses:
            # outfile.write(f"{json.dumps(response)}\n")
        # responses = []
    # print(len(responses))

# # except Exception as e:
    # # print(f"TPR: {TP/(TP+FN)}\nFPR: {TN/(TN+FP)} \nAcc: {(TN+TP)/(TP+FP+TN+FN)}  ")
# print(len(responses))