import torchfrom transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamerfrom threading import Thread model_id = "surajpuku/Vaan-V1_400M" # Load tokenizer + modeltokenizer = AutoTokenizer.from_pretrained(model_id) model = AutoModelForCausalLM.from_pretrained( model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, device_map="auto" if torch.cuda.is_available() else None) if not torch.cuda.is_available(): model.to("cpu") def stream_generate(prompt): inputs = tokenizer(prompt, return_tensors="pt").to(model.device) streamer = TextIteratorStreamer( tokenizer, skip_prompt=True, skip_special_tokens=True ) generation_kwargs = dict( **inputs, max_new_tokens=256, temperature=0.7, top_p=0.9, do_sample=True, streamer=streamer ) thread = Thread(target=model.generate, kwargs=generation_kwargs) thread.start() for token in streamer: print(token, end="", flush=True) print() print("Chat started. Type 'exit' to quit.") while True: user_input = input("\nYou: ") if user_input.lower() in ["exit", "quit"]: break print("Assistant: ", end="", flush=True) stream_generate(user_input)