ChaeSJ/llama-3.1-8b-finetuned

Code to download files from Hugging Face

import os
from huggingface_hub import snapshot_download

# 경로 설정
cache_dir = "--Enter Your Desired File Direction--"
tmp_dir   = "--Enter Your Desired File Direction--"
local_dir = "--Enter Your Desired File Direction--"
repo_id   = "ChaeSJ/llama-3.1-8b-finetuned"

os.environ["TRANSFORMERS_CACHE"] = cache_dir
os.environ["HF_HOME"] = cache_dir
os.makedirs(cache_dir, exist_ok=True)
os.makedirs(tmp_dir, exist_ok=True)
os.environ["TMPDIR"] = tmp_dir

print(f"Downloading snapshot of {repo_id} ...")
snapshot_download(
    repo_id=repo_id,
    local_dir=local_dir,
    local_dir_use_symlinks=False,
    cache_dir=cache_dir,
    resume_download=True
)
print(f" Downloaded all files to: {local_dir}")

Code that allows you to load the model directly from Huggingface without downloading it separately and proceed to code generation

import os
import json
import re
import torch
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForCausalLM

REPO_ID   = "ChaeSJ/llama-3.1-8b-finetuned"
SUBFOLDER = "llama_3_1_8b_finetuned"   
CACHE_DIR = "--Enter Your Desired File Direction--"   
TMP_DIR   = "--Enter Your Desired File Direction--"       

os.environ["TRANSFORMERS_VERBOSITY"] = "error"
os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
os.environ["HF_HOME"] = CACHE_DIR
os.makedirs(CACHE_DIR, exist_ok=True)
os.makedirs(TMP_DIR, exist_ok=True)
os.environ["TMPDIR"] = TMP_DIR

## 로컬 입출력 경로(프롬프트/결과물)
input_path       = "--Enter Your Desired File Direction--/selected_prompts_c++.json"
json_output_path = "--Enter Your Desired File Direction--/llama_3_1_8b_demo_generated.json"
cpp_output_dir   = "--Enter Your Desired File Direction--/extracted_cpp"

## 코드 블록에서 C/C++ 추출
def extract_cpp_code(text: str) -> str:
    m = re.findall(r"```(?:cpp|c\+\+)\n(.*?)```", text, re.DOTALL | re.IGNORECASE)
    if m:
        return m[0].strip()
    m = re.findall(r"```c\n(.*?)```", text, re.DOTALL | re.IGNORECASE)
    if m:
        return m[0].strip()
    return text.strip()

def apply_template(tokenizer, prompt_text: str) -> str:
    return tokenizer.apply_chat_template(
        [{"role": "user", "content": prompt_text}],
        tokenize=False,
        add_generation_prompt=True,
    )

## 토크나이저 & 모델: 허깅페이스에서 직접 로드
print(">> Loading tokenizer from Hugging Face...")
tokenizer = AutoTokenizer.from_pretrained(
    REPO_ID,
    subfolder=SUBFOLDER,
    cache_dir=CACHE_DIR,
    use_fast=True,
    padding_side="right",
)

if tokenizer.pad_token is None:
    tokenizer.add_special_tokens({"pad_token": tokenizer.eos_token})
    print(">> Added pad_token as eos_token")

print(">> Loading model from Hugging Face (full finetuned weights)...")
model = AutoModelForCausalLM.from_pretrained(
    REPO_ID,
    subfolder=SUBFOLDER,
    cache_dir=CACHE_DIR,
    torch_dtype=torch.bfloat16,  
    device_map="auto",           
    low_cpu_mem_usage=True,
)

if model.config.pad_token_id is None:
    model.config.pad_token_id = tokenizer.pad_token_id

model.eval()

## 프롬프트 로드
print(">> Loading prompts...")
with open(input_path, "r", encoding="utf-8") as f:
    prompts_data = json.load(f)

## 생성 & 저장
print(">> Generating code...")
os.makedirs(os.path.dirname(json_output_path), exist_ok=True)
os.makedirs(cpp_output_dir, exist_ok=True)

results = []

@torch.inference_mode()
def generate_one(prompt_text: str, max_input_len=1536, max_new_tokens=512) -> str:
    templated = apply_template(tokenizer, prompt_text)

    inputs = tokenizer(
        templated,
        return_tensors="pt",
        padding=False,
        truncation=True,
        max_length=max_input_len,
    )

    for k in inputs:
        inputs[k] = inputs[k].to(model.device)

    out = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,

        do_sample=False,
        num_beams=1,
        temperature=0.0,
        top_p=1.0,
        top_k=0,
        repetition_penalty=1.0,
        no_repeat_ngram_size=0,

        use_cache=True,
        return_dict_in_generate=True,
    )

    input_len = inputs["input_ids"].shape[1]
    gen_ids = out.sequences[0][input_len:]
    return tokenizer.decode(gen_ids, skip_special_tokens=True)

for idx, item in enumerate(tqdm(prompts_data, desc="Generating")):
    prompt_text = item.get("nl_prompt", "")
    prompt_id   = str(item.get("Prompt ID", "")).strip() or None

    generated_text = generate_one(prompt_text, max_input_len=1536, max_new_tokens=512)

    results.append({
        "Prompt ID": prompt_id if prompt_id is not None else f"{idx+1}",
        "nl_prompt": prompt_text,
        "generated_code": generated_text
    })

## JSON 저장
print(f">> Saving generated results to {json_output_path}")
with open(json_output_path, "w", encoding="utf-8") as f:
    json.dump(results, f, indent=2, ensure_ascii=False)

## C/C++ 코드 추출 및 파일 저장
print(f">> Extracting cpp code to directory: {cpp_output_dir}")
saved = 0
for i, item in enumerate(results):
    raw_code = item.get("generated_code", "")
    code = extract_cpp_code(raw_code)

    pid = item.get("Prompt ID")
    if pid and pid != "unknown":
        safe_pid = re.sub(r"[^a-zA-Z0-9_\-\.]+", "_", str(pid))[:64]
        filename = f"Llama_3.1_8b_demo_{safe_pid}.cpp"
    else:
        filename = f"Llama_3.1_8b_demo_{i+1:04d}.cpp"

    filepath = os.path.join(cpp_output_dir, filename)
    with open(filepath, "w", encoding="utf-8") as w:
        w.write(code)
    saved += 1

print(f">> Extraction complete. {saved} files saved to {cpp_output_dir}")