Views
No views yet
pip install transformers==4.49.0 torch==2.6.01from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
2import torch
3
4model_path = "Felladrin/Qwen2-96M"
5prompt = "I've been thinking about"
6device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7tokenizer = AutoTokenizer.from_pretrained(model_path)
8model = AutoModelForCausalLM.from_pretrained(model_path).to(device)
9streamer = TextStreamer(tokenizer)
10inputs = tokenizer(prompt, return_tensors="pt").to(device)
11model.generate(
12 inputs.input_ids,
13 attention_mask=inputs.attention_mask,
14 max_length=tokenizer.model_max_length,
15 streamer=streamer,
16 eos_token_id=tokenizer.eos_token_id,
17 pad_token_id=tokenizer.pad_token_id,
18 repetition_penalty=1.1,
19 do_sample=True,
20 temperature=0.7,
21 top_p=0.9,
22 top_k=0,
23 min_p=0.1,
24)