Views
No views yet
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4model = AutoModelForCausalLM.from_pretrained("OEvortex/BabyMistral")
5tokenizer = AutoTokenizer.from_pretrained("OEvortex/BabyMistral")
6
7# Define the chat input
8chat = [
9# { "role": "system", "content": "You are BabyMistral" },
10 { "role": "user", "content": "Hey there! How are you? 😊" }
11]
12
13inputs = tokenizer.apply_chat_template(
14 chat,
15 add_generation_prompt=True,
16 return_tensors="pt"
17).to(model.device)
18
19
20# Generate text
21outputs = model.generate(
22 inputs,
23 max_new_tokens=256,
24 do_sample=True,
25 temperature=0.6,
26 top_p=0.9,
27 eos_token_id=tokenizer.eos_token_id,
28
29
30)
31
32response = outputs[0][inputs.shape[-1]:]
33print(tokenizer.decode(response, skip_special_tokens=True))
34
35#I am doing well! How can I assist you today? 😊
36