107 lines
3.9 KiB
Python
107 lines
3.9 KiB
Python
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!")
|