#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 = "" 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 = "" 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) 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"[INST] {overarching_instruction}\nProgram 1: {prog1}\nProgram 2: {prog2} [/INST] {output_text} \n" prompt += f"[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]>" 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 = "" # 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 = "" # 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) # 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"[INST] {overarching_instruction}\nProgram 1: {prog1}\nProgram 2: {prog2} [/INST] {output_text} \n" # prompt += f"[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]>" # 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))