from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamerimport torchfrom threading import Thread model = AutoModelForCausalLM.from_pretrained("Aquiles-ai/Athenea-4B-Thinking", dtype=torch.bfloat16, trust_remote_code=True, device_map="auto", attn_implementation="flash_attention_2") tokenizer = AutoTokenizer.from_pretrained("Aquiles-ai/Athenea-4B-Thinking", trust_remote_code=True) messages = [ {"role": "user", "content": "Hey, explain the difference between artificial intelligence, machine learning, and deep learning."}] inputs = tokenizer.apply_chat_template( messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt",).to('cuda') # Create the streamerstreamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) # Build kwargs for generategenerate_kwargs = dict( **inputs, max_new_tokens=8092, pad_token_id=tokenizer.eos_token_id, eos_token_id=tokenizer.eos_token_id, streamer=streamer,) def _generate_thread(model, kwargs): with torch.no_grad(): model.generate(**kwargs) thread = Thread(target=_generate_thread, args=(model, generate_kwargs)) thread.start() for chunk in streamer: print(chunk, end="", flush=True)