DOUDOU_146_21 / compare_generation.py

Upload folder using huggingface_hub

77256a5 verified 11 months ago

3.86 kB

#!/usr/bin/env python3

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load model and tokenizer (same as server.py)
model_name = "models/Llama-3.2-1B-Instruct"
tok = AutoTokenizer.from_pretrained(model_name)
lm = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="cuda",
).eval()

def chat_current(system_prompt: str, user_prompt: str) -> str:
    """
    Current implementation (same as server.py) - will show warnings
    """
    print("🔴 Running CURRENT implementation (with warnings)...")
    
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]

    input_ids = tok.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(lm.device)

    with torch.inference_mode():
        output_ids = lm.generate(
            input_ids,  # No attention_mask, no pad_token_id
            max_new_tokens=2048,
            do_sample=True,
            temperature=0.2,
            repetition_penalty=1.1,
            top_k=100,
            top_p=0.95,
        )

    answer = tok.decode(
        output_ids[0][input_ids.shape[-1]:],
        skip_special_tokens=True,
        clean_up_tokenization_spaces=True,
    )
    return answer.strip()


def chat_fixed(system_prompt: str, user_prompt: str) -> str:
    """
    Fixed implementation - proper attention mask and pad token
    """
    print("🟢 Running FIXED implementation (no warnings)...")
    
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]

    # Get both input_ids and attention_mask
    inputs = tok.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True  # Returns dict with input_ids and attention_mask
    )
    
    # Move to device
    input_ids = inputs["input_ids"].to(lm.device)
    attention_mask = inputs["attention_mask"].to(lm.device)

    with torch.inference_mode():
        output_ids = lm.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,  # Proper attention mask
            pad_token_id=tok.eos_token_id,  # Explicit pad token
            max_new_tokens=2048,
            do_sample=True,
            temperature=0.2,
            repetition_penalty=1.1,
            top_k=100,
            top_p=0.95,
        )

    answer = tok.decode(
        output_ids[0][input_ids.shape[-1]:],
        skip_special_tokens=True,
        clean_up_tokenization_spaces=True,
    )
    return answer.strip()


def compare_generations():
    """Compare both implementations"""
    system_prompt = "You are a helpful assistant who tries to help answer the user's question."
    user_prompt = "Create a report on anxiety in work. How do I manage time and stress effectively?"
    
    print("=" * 60)
    print("COMPARING GENERATION METHODS")
    print("=" * 60)
    print(f"System: {system_prompt}")
    print(f"User: {user_prompt}")
    print("=" * 60)
    
    # Test current implementation
    print("\n" + "=" * 60)
    current_output = chat_current(system_prompt, user_prompt)
    print(f"CURRENT OUTPUT:\n{current_output}")
    
    print("\n" + "=" * 60)
    # Test fixed implementation  
    fixed_output = chat_fixed(system_prompt, user_prompt)
    print(f"FIXED OUTPUT:\n{fixed_output}")
    
    print("\n" + "=" * 60)
    print("COMPARISON:")
    print(f"Outputs are identical: {current_output == fixed_output}")
    print(f"Current length: {len(current_output)} chars")
    print(f"Fixed length: {len(fixed_output)} chars")


if __name__ == "__main__":
    # Set pad token for the fixed version
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    
    compare_generations()

	#!/usr/bin/env python3

	import torch
	from transformers import AutoModelForCausalLM, AutoTokenizer

	# Load model and tokenizer (same as server.py)
	model_name = "models/Llama-3.2-1B-Instruct"
	tok = AutoTokenizer.from_pretrained(model_name)
	lm = AutoModelForCausalLM.from_pretrained(
	model_name,
	torch_dtype=torch.bfloat16,
	device_map="cuda",
	).eval()

	def chat_current(system_prompt: str, user_prompt: str) -> str:
	"""
	Current implementation (same as server.py) - will show warnings
	"""
	print("🔴 Running CURRENT implementation (with warnings)...")

	messages = [
	{"role": "system", "content": system_prompt},
	{"role": "user", "content": user_prompt},
	]

	input_ids = tok.apply_chat_template(
	messages,
	add_generation_prompt=True,
	return_tensors="pt"
	).to(lm.device)

	with torch.inference_mode():
	output_ids = lm.generate(
	input_ids, # No attention_mask, no pad_token_id
	max_new_tokens=2048,
	do_sample=True,
	temperature=0.2,
	repetition_penalty=1.1,
	top_k=100,
	top_p=0.95,
	)

	answer = tok.decode(
	output_ids[0][input_ids.shape[-1]:],
	skip_special_tokens=True,
	clean_up_tokenization_spaces=True,
	)
	return answer.strip()


	def chat_fixed(system_prompt: str, user_prompt: str) -> str:
	"""
	Fixed implementation - proper attention mask and pad token
	"""
	print("🟢 Running FIXED implementation (no warnings)...")

	messages = [
	{"role": "system", "content": system_prompt},
	{"role": "user", "content": user_prompt},
	]

	# Get both input_ids and attention_mask
	inputs = tok.apply_chat_template(
	messages,
	add_generation_prompt=True,
	return_tensors="pt",
	return_dict=True # Returns dict with input_ids and attention_mask
	)

	# Move to device
	input_ids = inputs["input_ids"].to(lm.device)
	attention_mask = inputs["attention_mask"].to(lm.device)

	with torch.inference_mode():
	output_ids = lm.generate(
	input_ids=input_ids,
	attention_mask=attention_mask, # Proper attention mask
	pad_token_id=tok.eos_token_id, # Explicit pad token
	max_new_tokens=2048,
	do_sample=True,
	temperature=0.2,
	repetition_penalty=1.1,
	top_k=100,
	top_p=0.95,
	)

	answer = tok.decode(
	output_ids[0][input_ids.shape[-1]:],
	skip_special_tokens=True,
	clean_up_tokenization_spaces=True,
	)
	return answer.strip()


	def compare_generations():
	"""Compare both implementations"""
	system_prompt = "You are a helpful assistant who tries to help answer the user's question."
	user_prompt = "Create a report on anxiety in work. How do I manage time and stress effectively?"

	print("=" * 60)
	print("COMPARING GENERATION METHODS")
	print("=" * 60)
	print(f"System: {system_prompt}")
	print(f"User: {user_prompt}")
	print("=" * 60)

	# Test current implementation
	print("\n" + "=" * 60)
	current_output = chat_current(system_prompt, user_prompt)
	print(f"CURRENT OUTPUT:\n{current_output}")

	print("\n" + "=" * 60)
	# Test fixed implementation
	fixed_output = chat_fixed(system_prompt, user_prompt)
	print(f"FIXED OUTPUT:\n{fixed_output}")

	print("\n" + "=" * 60)
	print("COMPARISON:")
	print(f"Outputs are identical: {current_output == fixed_output}")
	print(f"Current length: {len(current_output)} chars")
	print(f"Fixed length: {len(fixed_output)} chars")


	if __name__ == "__main__":
	# Set pad token for the fixed version
	if tok.pad_token is None:
	tok.pad_token = tok.eos_token

	compare_generations()