Create test_prompts.py

test_prompts.py

@@ -0,0 +1,106 @@
+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!")