Update test_prompts.py

test_prompts.py

@@ -4,7 +4,7 @@ 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"
+CHECKPOINT = "HuggingFaceTB/smollm-350M-instruct-add-basics"
 
 print(f"💾 Loading the model and tokenizer: {CHECKPOINT}...")
 device = "cuda"
@@ -13,21 +13,56 @@ model_s = AutoModelForCausalLM.from_pretrained(CHECKPOINT).to(device)
 
 print("🧪 Testing single-turn conversations...")
 L = [
-    "Hi",
-    "Hello",
+    # Witing and general knowledge prompts
+    "Discuss the ethical implications of using AI in hiring processes.",
+    "Give me some tips to improve my time management skills?",
+    "Write a short dialogue between a customer and a waiter at a restaurant.",
+    "wassup?",
     "Tell me a joke",
-    "Who are you?",
-    "What's your name?",
+    "Hi, what are some popular dishes from Japan?",
+    "What is the capital of Switzerland?",
+    "What is the capital of France?",
+    "What's the capital of Portugal?",
+    "What is the capital of Morocco?",
     "How do I make pancakes?",
+    "Write a poem about Helium",
+    "Do you think it's important for a company to have a strong company culture? Why or why not?",
+    "What is your favorite book?",
+    "What is the most interesting fact you know?",
+    "What is your favorite movie?",
+
+    # Science prompts
     "Can you tell me what is gravity?",
-    "What is the capital of Morocco?",
+    "Who discovered gravity?",
+    "How does a rainbow form?",
+    "What are the three states of matter?",
+    "Why is the sky blue?",
+    "What is the water cycle?",
+    "How do magnets work?",
+    "What is buoyancy?",
+    "What is the speed of light?",
     "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?",
-]
+    "what's the sum of 2 and 2?",
+    "what's the sum of 2 and 3?",
+    "What is the term for the process by which plants make their own food?",
+    "If you have 8 apples and you give away 3, how many apples do you have left?",
 
+    # Python prompts
+    "How do I define a function in Python?",
+    "Can you explain what a dictionary is in Python?",
+    "Write a sort alrogithm in Python",
+    "Write a fibonacci sequence in Python",
+    "How do I read a file in Python?",
+    "How do I make everything uppercase in Python?",
+    "implement bubble sort in Python",
+
+    # Creative prompts
+    "Write a short story about a time traveler",
+    "Describe a futuristic city in three sentences",
+    "Describe a new color that doesn't exist",
+    "Create a slogan for a time machine company",
+    "Describe a world where plants can speak",
+]
 
 for i in range(len(L)):
     print(f"🔮 {L[i]}")
@@ -44,63 +79,3 @@ for i in range(len(L)):
         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!")