falcon_inf.py

import os
import json
import torch
import logging
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForCausalLM
from prompt import sp1, sp2

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

INPUT_DIR = "data"
RESULT_DIR = "results"
ERROR_DIR = "error"

os.makedirs(RESULT_DIR, exist_ok=True)
os.makedirs(ERROR_DIR, exist_ok=True)

MODEL_NAME = "tiiuae/falcon-7b-instruct"
FILE_NAME = "java_xl.json"
QUERY = "prompt" #sp1 or sp2

logging.info(f"Loading model: {MODEL_NAME}")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map="auto",
)

def get_valid_files(directory):
    files = sorted(os.listdir(directory))
    return [file for file in files if not file.startswith(('.', '~'))]

def load_data(file_path):
    try:
        with open(file_path, "r") as f:
            return json.load(f)
    except Exception as e:
        logging.error(f"Error loading file {file_path}: {e}")
        return []

def generate_response(codeA, codeB):
    input_text = f">>QUESTION<<{QUERY}\nCode A:\n{codeA}\nCode B:\n{codeB}\n>>ANSWER<<"
    input_ids = tokenizer.encode(input_text, return_tensors="pt").to("cuda")
    
    output = model.generate(
        input_ids,
        max_length=4000,
        do_sample=True,
        top_k=10,
        num_return_sequences=1,
        eos_token_id=tokenizer.eos_token_id,
    )

    return tokenizer.decode(output[0], skip_special_tokens=False)

def process_file(input_path, output_path, error_path):
    data = load_data(input_path)
    
    for element in tqdm(data, desc=f"Processing {os.path.basename(input_path)}"):
        try:
            element["result"] = generate_response(element['codeA'], element['codeB'])
        except Exception as e:
            logging.error(f"Error processing element: {e}")
            with open(error_path, "w") as f:
                json.dump(data, f, indent=4)

    with open(output_path, "w") as f:
        json.dump(data, f, indent=4)

def main():
    files = get_valid_files(INPUT_DIR)
    for file in files:
        file_path = os.path.join(INPUT_DIR, file)
        output_path = os.path.join(RESULT_DIR, f"llma_2_{file}")
        error_path = os.path.join(ERROR_DIR, f"llma_error_{file}")

        logging.info(f"Processing file: {file}")
        process_file(file_path, output_path, error_path)

if __name__ == "__main__":
    main()