LLM → narrative summary
Code → structured memory update

completion_only_loss=True

packing=False

[PREVIOUS MEMORY]
User is designing a multi-turn RAG system

[RECENT TURNS]
User: Should memory be separate?
Assistant: It can be separate
User: What data can we use?

[ACTIVE DOCUMENT]
Survey of multi-turn LLM systems

[PREVIOUS MEMORY]
[RECENT TURNS]
[ACTIVE DOCUMENT]

# ============================================================
# 1. base model 추가 로드
#    이미 fine-tuned model은 model 변수에 있다고 가정
# ============================================================
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import time
import json

BASE_MODEL_ID = "Qwen/Qwen2.5-3B-Instruct"

base_tokenizer = tokenizer  # 이미 같은 계열이면 기존 tokenizer 재사용 가능

base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_ID,
    torch_dtype="auto",
    device_map="auto",
    trust_remote_code=True
)

base_tokenizer.pad_token = base_tokenizer.eos_token
base_tokenizer.padding_side = "right"

base_model.config.pad_token_id = base_tokenizer.pad_token_id
base_model.config.eos_token_id = base_tokenizer.eos_token_id

if hasattr(base_model, "generation_config") and base_model.generation_config is not None:
    base_model.generation_config.pad_token_id = base_tokenizer.pad_token_id
    base_model.generation_config.eos_token_id = base_tokenizer.eos_token_id

print("base model loaded")

# ============================================================
# 2. 비교용 입력
# ============================================================
sample_previous_memory = {
    "narrative": "The user is designing a multi-turn RAG system.",
    "structured": {
        "goal": "Design a multi-turn architecture",
        "established_facts": ["RAG is used"],
        "current_focus": "memory model design",
        "unresolved_questions": ["whether to separate the memory model service"],
    },
}

sample_recent_turns = [
    {"role": "user", "text": "Should the memory model run as a separate endpoint?"},
    {"role": "assistant", "text": "It can run on the same machine but as a separate service."},
    {"role": "user", "text": "Then what public datasets can we use for training?"},
]

sample_active_doc = {
    "file_name": "multiturn_survey.pdf",
    "doc_summary": "A survey of multi-turn interactions with large language models."
}

def render_recent_turns(turns):
    return "\n".join([f'{t["role"].capitalize()}: {t["text"]}' for t in turns])

memory_prompt_system = (
    "You are a session-memory summarization model. "
    "Summarize only durable and useful information for future turns. "
    "Always include the user's current focus and latest intent when important. "
    "Do not invent facts."
)

memory_prompt_user = f"""
Create a concise memory summary for future dialogue turns.

[PREVIOUS MEMORY]
{json.dumps(sample_previous_memory, ensure_ascii=False, indent=2)}

[RECENT TURNS]
{render_recent_turns(sample_recent_turns)}

[ACTIVE DOCUMENT]
{json.dumps(sample_active_doc, ensure_ascii=False, indent=2)}

Return a concise summary in plain English.
"""

messages = [
    {"role": "system", "content": memory_prompt_system},
    {"role": "user", "content": memory_prompt_user},
]

prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)

inputs = tokenizer(
    prompt,
    return_tensors="pt",
    truncation=True,
    max_length=3072
)

# ============================================================
# 3. 공통 generate 함수
# ============================================================
def run_generate(model_obj, tokenizer_obj, inputs_dict, max_new_tokens=160):
    model_inputs = {k: v.to(model_obj.device) for k, v in inputs_dict.items()}

    if torch.cuda.is_available():
        torch.cuda.synchronize()
    start = time.perf_counter()

    with torch.no_grad():
        outputs = model_obj.generate(
            **model_inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            temperature=0.0,
            eos_token_id=tokenizer_obj.eos_token_id,
            pad_token_id=tokenizer_obj.pad_token_id,
        )

    if torch.cuda.is_available():
        torch.cuda.synchronize()
    end = time.perf_counter()

    gen_ids = outputs[0][model_inputs["input_ids"].shape[1]:]
    gen_text = tokenizer_obj.decode(gen_ids, skip_special_tokens=True).strip()

    return {
        "text": gen_text,
        "elapsed": end - start,
        "output_tokens": len(gen_ids),
        "tokens_per_sec": (len(gen_ids) / (end - start)) if (end - start) > 0 else None,
    }

# ============================================================
# 4. fine-tuned vs base 비교
#    model = 네가 이미 로드한 fine-tuned model
# ============================================================
ft_result = run_generate(model, tokenizer, inputs, max_new_tokens=160)
base_result = run_generate(base_model, base_tokenizer, inputs, max_new_tokens=160)

print("============== FINE-TUNED MODEL ==============")
print(ft_result["text"])
print(f"time: {ft_result['elapsed']:.4f}s | tokens: {ft_result['output_tokens']} | tok/s: {ft_result['tokens_per_sec']:.2f}")

print("\n============== BASE MODEL ==============")
print(base_result["text"])
print(f"time: {base_result['elapsed']:.4f}s | tokens: {base_result['output_tokens']} | tok/s: {base_result['tokens_per_sec']:.2f}")