File size: 3,980 Bytes
d26da8b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
from transformers import AutoModelForCausalLM, AutoTokenizer
BASE_PATH = "/fsx/loubna/projects/alignment-handbook/recipes/cosmo2/sft/data"
TEMPERATURE = 0.2
TOP_P = 0.9
CHECKPOINT = "loubnabnl/smollm-350M-instruct-add-basics"
print(f"💾 Loading the model and tokenizer: {CHECKPOINT}...")
device = "cuda"
tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
model_s = AutoModelForCausalLM.from_pretrained(CHECKPOINT).to(device)
print("🧪 Testing single-turn conversations...")
L = [
"Hi",
"Hello",
"Tell me a joke",
"Who are you?",
"What's your name?",
"How do I make pancakes?",
"Can you tell me what is gravity?",
"What is the capital of Morocco?",
"What's 2+2?",
"Hi, what is 2+1?",
"What's 3+5?",
"Write a poem about Helium",
"Hi, what are some popular dishes from Japan?",
]
for i in range(len(L)):
print(f"🔮 {L[i]}")
messages = [{"role": "user", "content": L[i]}]
input_text = tokenizer.apply_chat_template(messages, tokenize=False)
inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
outputs = model_s.generate(
inputs, max_new_tokens=200, top_p=TOP_P, do_sample=True, temperature=TEMPERATURE
)
with open(
f"{BASE_PATH}/{CHECKPOINT.split('/')[-1]}_temp_{TEMPERATURE}_topp{TOP_P}.txt",
"a",
) as f:
f.write("=" * 50 + "\n")
f.write(tokenizer.decode(outputs[0]))
f.write("\n")
print("🧪 Now testing multi-turn conversations...")
# Multi-turn conversations
messages_1 = [
{"role": "user", "content": "Hi"},
{"role": "assistant", "content": "Hello! How can I help you today?"},
{"role": "user", "content": "What's 2+2?"},
]
messages_2 = [
{"role": "user", "content": "Hi"},
{"role": "assistant", "content": "Hello! How can I help you today?"},
{"role": "user", "content": "What's 2+2?"},
{"role": "assistant", "content": "4"},
{"role": "user", "content": "Why?"},
]
messages_3 = [
{"role": "user", "content": "Who are you?"},
{"role": "assistant", "content": "I am an AI assistant. How can I help you today?"},
{"role": "user", "content": "What's your name?"},
]
messages_4 = [
{"role": "user", "content": "Tell me a joke"},
{"role": "assistant", "content": "Sure! Why did the tomato turn red?"},
{"role": "user", "content": "Why?"},
]
messages_5 = [
{"role": "user", "content": "Can you tell me what is gravity?"},
{
"role": "assistant",
"content": "Sure! Gravity is a force that attracts objects toward each other. It is what keeps us on the ground and what makes things fall.",
},
{"role": "user", "content": "Who discovered it?"},
]
messages_6 = [
{"role": "user", "content": "How do I make pancakes?"},
{
"role": "assistant",
"content": "Sure! Here is a simple recipe for pancakes: Ingredients: 1 cup flour, 1 cup milk, 1 egg, 1 tbsp sugar, 1 tsp baking powder, 1/2 tsp salt. Instructions: 1. Mix all the dry ingredients together in a bowl. 2. Add the milk and egg and mix until smooth. 3. Heat a non-stick pan over medium heat. 4. Pour 1/4 cup of batter onto the pan. 5. Cook until bubbles form on the surface, then flip and cook for another minute. 6. Serve with your favorite toppings.",
},
{"role": "user", "content": "What are some popular toppings?"},
]
L = [messages_1, messages_2, messages_3, messages_4, messages_5, messages_6]
for i in range(len(L)):
input_text = tokenizer.apply_chat_template(L[i], tokenize=False)
inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
outputs = model_s.generate(
inputs, max_new_tokens=200, top_p=TOP_P, do_sample=True, temperature=TEMPERATURE
)
with open(
f"{BASE_PATH}/{CHECKPOINT.split('/')[-1]}_temp_{TEMPERATURE}_topp{TOP_P}_MT.txt",
"a",
) as f:
f.write("=" * 50 + "\n")
f.write(tokenizer.decode(outputs[0]))
f.write("\n")
print("🔥 Done!")
|