llama2_inf.py

import os
import json
import torch
import logging
import transformers
from tqdm import tqdm
from langchain_huggingface import HuggingFacePipeline
from transformers import AutoTokenizer
from langchain.prompts import PromptTemplate

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 = "meta-llama/Llama-2-7b-chat-hf"

logging.info(f"Loading model: {MODEL_NAME}")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
pipeline = transformers.pipeline(
    "text-generation",
    model=MODEL_NAME,
    tokenizer=tokenizer,
    torch_dtype=torch.float16,
    trust_remote_code=True,
    device_map="auto",
    max_length=5000,
    do_sample=True,
    truncation=True,
    top_k=10,
    top_p=0.9,
    num_return_sequences=1,
    repetition_penalty=1.1,
    eos_token_id=tokenizer.eos_token_id,
)

llm = HuggingFacePipeline(pipeline=pipeline, model_kwargs={'temperature': 0.3})

TEMPLATE = """[INST]
Consider the overall structure and logic of the following two code snippets and determine if they perform a similar task. 
Respond with 'yes' if they perform similar tasks or 'no' otherwise.

Code A:
{codeA}

Code B:
{codeB}
[/INST]"""
prompt = PromptTemplate(template=TEMPLATE, input_variables=["codeA", "codeB"])
llm_chain = prompt | llm

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

def process_file(file_path, output_path, error_path):
    try:
        with open(file_path, "r") as f:
            data = json.load(f)

        for element in tqdm(data, desc=f"Processing {os.path.basename(file_path)}"):
            inputs = {"codeA": element['codeA'], "codeB": element['codeB']}
            element["result"] = llm_chain.invoke(inputs)

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

    except Exception as e:
        logging.error(f"Error processing {file_path}: {e}")
        with open(error_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"llama2_{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()