from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

path = "spark_v4_fp16_final"
tokenizer = AutoTokenizer.from_pretrained(path)
model = AutoModelForCausalLM.from_pretrained(path).to("cuda")

prompts = [
    "Artificial Intelligence is",
    "The main concept of physics is",
    "In the year 1969, "
]

for prompt in prompts:
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(
        **inputs, 
        max_new_tokens=200, 
        do_sample=True, 
        top_k=25, 
        temperature=0.8,
        pad_token_id=tokenizer.eos_token_id
    )
    print(f"PROMPT: {prompt}")
    print(f"OUTPUT: {tokenizer.decode(outputs[0], skip_special_tokens=True)}\n{'-'*40}")