from dotenv import load_dotenv
load_dotenv()

import time
from transformers import pipeline, GenerationConfig, AutoTokenizer

# 1. Load your trained model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")

generator = pipeline(
    "text-generation",
    model="./trained_model",
    tokenizer=tokenizer,
    clean_up_tokenization_spaces=False,
)

# 2. Configure generation settings
generation_config = GenerationConfig(
    max_new_tokens=50,
    do_sample=True,
)

# 3. Give it a starting prompt
prompt = "what is 1+10"

# 4. Generate the response with timing
input_ids = tokenizer.encode(prompt, return_tensors="pt")
tokens_in = input_ids.shape[1]

start_time = time.time()
results = generator(prompt, generation_config=generation_config)
elapsed = time.time() - start_time

# 5. Calculate stats
generated_text = results[0]["generated_text"]
output_ids = tokenizer.encode(generated_text)
tokens_out = len(output_ids)
new_tokens = tokens_out - tokens_in
tokens_per_sec = new_tokens / elapsed if elapsed > 0 else 0

# 6. Print the output and stats
print("\n" + "=" * 60)
print("  GENERATED TEXT")
print("=" * 60)
print(generated_text)
print("=" * 60)
print(f"  📊 Stats:")
print(f"     Prompt tokens (in):   {tokens_in}")
print(f"     Output tokens (out):  {tokens_out}")
print(f"     New tokens generated: {new_tokens}")
print(f"     Generation time:      {elapsed:.2f}s")
print(f"     Speed:                {tokens_per_sec:.1f} tokens/sec")
print(f"     Model:                ./trained_model")
print(f"     Tokenizer:            Qwen/Qwen3-8B")
print(f"     Sampling:             {'yes' if generation_config.do_sample else 'no'}")
print(f"     Max new tokens:       {generation_config.max_new_tokens}")
print("=" * 60)