Upload full run contents (scripts, checkpoints, and results) (part 13)

babylm_eval_repo/strict/results/xpertgpt_fresh/main/zero_shot/causal/entity_tracking/entity_tracking/best_temperature_report.txt

@@ -0,0 +1,28 @@
+TEMPERATURE: 1.00
+
+### UID ACCURACY
+ambiref_0_ops: 17.72
+ambiref_1_ops: 17.29
+ambiref_2_ops: 18.40
+ambiref_3_ops: 19.80
+ambiref_4_ops: 22.12
+ambiref_5_ops: 25.20
+regular_0_ops: 18.38
+regular_1_ops: 15.89
+regular_2_ops: 23.21
+regular_3_ops: 22.12
+regular_4_ops: 22.16
+regular_5_ops: 27.66
+move_contents_0_ops: 17.44
+move_contents_1_ops: 19.68
+move_contents_2_ops: 18.30
+move_contents_3_ops: 17.98
+move_contents_4_ops: 20.11
+move_contents_5_ops: 25.86
+regular: 21.57
+ambiref: 20.09
+move_contents: 19.90
+
+### AVERAGE ACCURACY
+20.52
+

babylm_eval_repo/strict/results/xpertgpt_fresh/main/zero_shot/causal/ewok/ewok_filtered/best_temperature_report.txt

@@ -0,0 +1,37 @@
+TEMPERATURE: 1.00
+
+### UID ACCURACY
+social-interactions: 57.48
+physical-interactions: 48.92
+social-relations: 49.94
+material-dynamics: 47.01
+quantitative-properties: 60.19
+physical-relations: 51.34
+material-properties: 45.88
+agent-properties: 49.14
+physical-dynamics: 48.33
+spatial-relations: 48.78
+social-properties: 60.98
+
+### CONTEXT_TYPE ACCURACY
+direct: 50.03
+indirect: 50.76
+
+### CONTEXT_CONTRAST ACCURACY
+antonym: 51.81
+variable swap: 49.96
+other: 51.02
+variable_swap: 43.33
+negation: 48.42
+active-passive: 40.00
+material: 47.02
+number: 62.50
+game: 25.00
+
+### TARGET_CONTRAST ACCURACY
+concept swap: 51.07
+variable swap: 49.51
+
+### AVERAGE ACCURACY
+51.64
+

babylm_eval_repo/strict/results/xpertgpt_fresh/main/zero_shot/causal/global_piqa_nonparallel/global_piqa_nonparallel/best_temperature_report.txt

@@ -0,0 +1,107 @@
+TEMPERATURE: 1.00
+
+### UID ACCURACY
+group0123_ex000088_eng_latn_0_v1: 100.00
+group0123_ex000012_eng_latn_0_v1: 0.00
+group0123_ex000022_eng_latn_0_v1: 100.00
+group0123_ex000102_eng_latn_0_v1: 0.00
+group0123_ex000028_eng_latn_0_v1: 100.00
+group0123_ex000002_eng_latn_0_v1: 0.00
+group0123_ex000081_eng_latn_0_v1: 0.00
+group0123_ex000018_eng_latn_0_v1: 0.00
+group0123_ex000000_eng_latn_0_v1: 100.00
+group0123_ex000046_eng_latn_0_v1: 100.00
+group0123_ex000075_eng_latn_0_v1: 0.00
+group0123_ex000033_eng_latn_0_v1: 100.00
+group0123_ex000073_eng_latn_0_v1: 100.00
+group0123_ex000021_eng_latn_0_v1: 100.00
+group0123_ex000095_eng_latn_0_v1: 0.00
+group0123_ex000013_eng_latn_0_v1: 100.00
+group0123_ex000030_eng_latn_0_v1: 0.00
+group0123_ex000059_eng_latn_0_v1: 100.00
+group0123_ex000089_eng_latn_0_v1: 100.00
+group0123_ex000061_eng_latn_0_v1: 100.00
+group0123_ex000053_eng_latn_0_v1: 0.00
+group0123_ex000076_eng_latn_0_v1: 0.00
+group0123_ex000064_eng_latn_0_v1: 100.00
+group0123_ex000041_eng_latn_0_v1: 0.00
+group0123_ex000031_eng_latn_0_v1: 0.00
+group0123_ex000034_eng_latn_0_v1: 100.00
+group0123_ex000062_eng_latn_0_v1: 0.00
+group0123_ex000036_eng_latn_0_v1: 100.00
+group0123_ex000056_eng_latn_0_v1: 0.00
+group0123_ex000070_eng_latn_0_v1: 0.00
+group0123_ex000045_eng_latn_0_v1: 100.00
+group0123_ex000014_eng_latn_0_v1: 0.00
+group0123_ex000094_eng_latn_0_v1: 100.00
+group0123_ex000052_eng_latn_0_v1: 0.00
+group0123_ex000006_eng_latn_0_v1: 0.00
+group0123_ex000007_eng_latn_0_v1: 100.00
+group0123_ex000091_eng_latn_0_v1: 100.00
+group0123_ex000010_eng_latn_0_v1: 0.00
+group0123_ex000071_eng_latn_0_v1: 100.00
+group0123_ex000072_eng_latn_0_v1: 100.00
+group0123_ex000042_eng_latn_0_v1: 100.00
+group0123_ex000100_eng_latn_0_v1: 100.00
+group0123_ex000039_eng_latn_0_v1: 0.00
+group0123_ex000024_eng_latn_0_v1: 100.00
+group0123_ex000025_eng_latn_0_v1: 0.00
+group0123_ex000029_eng_latn_0_v1: 0.00
+group0123_ex000068_eng_latn_0_v1: 0.00
+group0123_ex000086_eng_latn_0_v1: 0.00
+group0123_ex000020_eng_latn_0_v1: 0.00
+group0123_ex000057_eng_latn_0_v1: 0.00
+group0123_ex000092_eng_latn_0_v1: 0.00
+group0123_ex000037_eng_latn_0_v1: 100.00
+group0123_ex000044_eng_latn_0_v1: 0.00
+group0123_ex000083_eng_latn_0_v1: 100.00
+group0123_ex000008_eng_latn_0_v1: 100.00
+group0123_ex000032_eng_latn_0_v1: 0.00
+group0123_ex000096_eng_latn_0_v1: 100.00
+group0123_ex000084_eng_latn_0_v1: 0.00
+group0123_ex000038_eng_latn_0_v1: 100.00
+group0123_ex000005_eng_latn_0_v1: 100.00
+group0123_ex000058_eng_latn_0_v1: 0.00
+group0123_ex000001_eng_latn_0_v1: 0.00
+group0123_ex000047_eng_latn_0_v1: 100.00
+group0123_ex000069_eng_latn_0_v1: 0.00
+group0123_ex000023_eng_latn_0_v1: 0.00
+group0123_ex000098_eng_latn_0_v1: 100.00
+group0123_ex000011_eng_latn_0_v1: 100.00
+group0123_ex000009_eng_latn_0_v1: 0.00
+group0123_ex000055_eng_latn_0_v1: 0.00
+group0123_ex000027_eng_latn_0_v1: 0.00
+group0123_ex000051_eng_latn_0_v1: 0.00
+group0123_ex000004_eng_latn_0_v1: 0.00
+group0123_ex000080_eng_latn_0_v1: 0.00
+group0123_ex000054_eng_latn_0_v1: 0.00
+group0123_ex000016_eng_latn_0_v1: 100.00
+group0123_ex000066_eng_latn_0_v1: 0.00
+group0123_ex000079_eng_latn_0_v1: 0.00
+group0123_ex000074_eng_latn_0_v1: 0.00
+group0123_ex000090_eng_latn_0_v1: 0.00
+group0123_ex000099_eng_latn_0_v1: 100.00
+group0123_ex000060_eng_latn_0_v1: 100.00
+group0123_ex000085_eng_latn_0_v1: 100.00
+group0123_ex000043_eng_latn_0_v1: 100.00
+group0123_ex000050_eng_latn_0_v1: 0.00
+group0123_ex000077_eng_latn_0_v1: 100.00
+group0123_ex000078_eng_latn_0_v1: 100.00
+group0123_ex000097_eng_latn_0_v1: 100.00
+group0123_ex000067_eng_latn_0_v1: 100.00
+group0123_ex000035_eng_latn_0_v1: 0.00
+group0123_ex000048_eng_latn_0_v1: 0.00
+group0123_ex000040_eng_latn_0_v1: 100.00
+group0123_ex000003_eng_latn_0_v1: 0.00
+group0123_ex000101_eng_latn_0_v1: 0.00
+group0123_ex000082_eng_latn_0_v1: 0.00
+group0123_ex000017_eng_latn_0_v1: 0.00
+group0123_ex000015_eng_latn_0_v1: 0.00
+group0123_ex000063_eng_latn_0_v1: 0.00
+group0123_ex000026_eng_latn_0_v1: 0.00
+group0123_ex000103_eng_latn_0_v1: 0.00
+group0123_ex000049_eng_latn_0_v1: 100.00
+
+### AVERAGE ACCURACY
+44.00
+

babylm_eval_repo/strict/results/xpertgpt_fresh/main/zero_shot/causal/global_piqa_nonparallel/global_piqa_nonparallel/predictions.json

@@ -0,0 +1 @@
+{"group0123_ex000088_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000088_eng_latn_0_v1_0", "pred": " use a muffin tin for the various condiments."}]}, "group0123_ex000012_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000012_eng_latn_0_v1_0", "pred": " tape holiday treats to the wall in pattern of a tree."}]}, "group0123_ex000022_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000022_eng_latn_0_v1_0", "pred": " Add whipped cream, nuts, and chocolate on top of ice cream."}]}, "group0123_ex000102_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000102_eng_latn_0_v1_0", "pred": " cup."}]}, "group0123_ex000028_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000028_eng_latn_0_v1_0", "pred": " Tackle the opponent in their own end zone."}]}, "group0123_ex000002_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000002_eng_latn_0_v1_0", "pred": " use a knife to crimp the edges."}]}, "group0123_ex000081_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000081_eng_latn_0_v1_0", "pred": " bite hard into the candy with your mouth closed in the dark."}]}, "group0123_ex000018_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000018_eng_latn_0_v1_0", "pred": " use an electric knife with the blade pointed away from the bird, and slice thin slices of turkey breast."}]}, "group0123_ex000000_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000000_eng_latn_0_v1_0", "pred": " To make pie crust, you will need flour, sugar, butter, and ice water."}]}, "group0123_ex000046_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000046_eng_latn_0_v1_0", "pred": " Add your dry ingredients to the bowl and stir them together well with a spoon, then add wet ingredients and stir it all together until joins, then use your hands to mix the dough until it forms a solid ball."}]}, "group0123_ex000075_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000075_eng_latn_0_v1_0", "pred": " hit more baseballs than the opposing team."}]}, "group0123_ex000033_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000033_eng_latn_0_v1_0", "pred": " Toast 2 frozen whole-wheat waffles. Spread one with peanut butter and the other with jelly. Sandwich together."}]}, "group0123_ex000073_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000073_eng_latn_0_v1_0", "pred": " Place a banana cut down the middle into a bowl and then place three scoops of ice cream along the banana. Add chocolate syrup on top if you'd like, and then a cherry on top."}]}, "group0123_ex000021_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000021_eng_latn_0_v1_0", "pred": " apply lemon juice to your hair and then sit in the sun."}]}, "group0123_ex000095_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000095_eng_latn_0_v1_0", "pred": " place 2-3 marshmallows in your bag of brown sugar. Replace the marshmallows when they are soft."}]}, "group0123_ex000013_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000013_eng_latn_0_v1_0", "pred": " many people use a pinched grip like they are holding a pencil or throwing a dart. You have much more control if you hold razor in between the handle and razor. In most circumstances try to hold the handle between you middle and pointer finger while steadying the blade with your thumb. Getting behind the neck can be tricky and you may have to modify your grip."}]}, "group0123_ex000030_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000030_eng_latn_0_v1_0", "pred": " Let the plaster soften."}]}, "group0123_ex000059_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000059_eng_latn_0_v1_0", "pred": " place 12 strawberries in a freezer and leave overnight to freeze. Place frozen strawberries in blender with 1/4 cup milk, 1 tablespoon of sugar, 1 cup plain yogurt, and 1/2 teaspoon of vanilla extract. Blend until thick and pureed."}]}, "group0123_ex000089_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000089_eng_latn_0_v1_0", "pred": " tighten the bow then apply violin rosin for 10 strokes."}]}, "group0123_ex000061_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000061_eng_latn_0_v1_0", "pred": " cannot cover a mop."}]}, "group0123_ex000053_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000053_eng_latn_0_v1_0", "pred": " by wrapping brown sugar in parchment paper and put in oven at 300 degrees for a few minutes."}]}, "group0123_ex000076_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000076_eng_latn_0_v1_0", "pred": " Place the metal end into the latch."}]}, "group0123_ex000064_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000064_eng_latn_0_v1_0", "pred": " use the front camera."}]}, "group0123_ex000041_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000041_eng_latn_0_v1_0", "pred": " Mix 2 cups of soda with one cup of sugar in a quart jar. Punch holes in the lid and sprinkle around the outside of your house."}]}, "group0123_ex000031_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000031_eng_latn_0_v1_0", "pred": " spray water on pole to make it slippery."}]}, "group0123_ex000034_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000034_eng_latn_0_v1_0", "pred": " Cover the beverage with wet paper towels and place it in the freezer."}]}, "group0123_ex000062_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000062_eng_latn_0_v1_0", "pred": " You can use a new t-shirt."}]}, "group0123_ex000036_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000036_eng_latn_0_v1_0", "pred": " Put desired number of eggs in a sauce pan and add enough water to cover. Put on a high heat and let the water come to boil. Continue to cook for 4 minutes. Remove from heat and drain."}]}, "group0123_ex000056_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000056_eng_latn_0_v1_0", "pred": " Put the spinach into a bowl. Run water over it constantly, and use your hands and fingers to help wash off any particles."}]}, "group0123_ex000070_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000070_eng_latn_0_v1_0", "pred": " Pour white vinegar into a small spray bottle. Spray the vinegar evenly onto your stovetop. Then sprinkle baking soda evenly across your stovetop. The baking soda will react with the vinegar, creating a fizzy cleaner that will loosen up grease and grime. Next, put a few drops of dish soap into a bucket or your sink, and fill it with cold water. Dunk your towel into the soapy water, then wring it out lightly so that it\u2019s wet but not dripping. Spread the wet towel over the stovetop, covering the baking soda and vinegar. Let it sit for about 15 minutes. Remove the towel and wipe away the remaining baking soda. To finish off, give your stovetop another spray with the vinegar, and use a microfiber cloth to buff it to a sparkling shine."}]}, "group0123_ex000045_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000045_eng_latn_0_v1_0", "pred": " Put a candle wick in the center of a can of crisco, it will last up to 45 days."}]}, "group0123_ex000014_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000014_eng_latn_0_v1_0", "pred": " swat it with a fly swatter then wait until it lands."}]}, "group0123_ex000094_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000094_eng_latn_0_v1_0", "pred": " use a pencil to mark a line in the center of the board before cutting it."}]}, "group0123_ex000052_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000052_eng_latn_0_v1_0", "pred": " Fill the base with water."}]}, "group0123_ex000006_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000006_eng_latn_0_v1_0", "pred": " 3 cups of strawberries, 3 cups of peaches, 2 cups of honey, and 2/3 cup of vanilla greek yogurt."}]}, "group0123_ex000007_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000007_eng_latn_0_v1_0", "pred": " add baking soda to the water while boiling."}]}, "group0123_ex000091_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000091_eng_latn_0_v1_0", "pred": " Place one foot in front of the other."}]}, "group0123_ex000010_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000010_eng_latn_0_v1_0", "pred": " Turn off the light switch and wait some time for the bulb to cool down if necessary. Place your fingers around the wide part of the bulb, and twist clockwise until the bulb comes out."}]}, "group0123_ex000071_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000071_eng_latn_0_v1_0", "pred": " turn the tip of spray nozzle to \"spray\" then pump the trigger until spray starts coming out."}]}, "group0123_ex000072_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000072_eng_latn_0_v1_0", "pred": " Place cardboard inside the windshield."}]}, "group0123_ex000042_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000042_eng_latn_0_v1_0", "pred": " Carefully cut off any areas which are growing too quickly."}]}, "group0123_ex000100_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000100_eng_latn_0_v1_0", "pred": " can tighten a nut easily."}]}, "group0123_ex000039_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000039_eng_latn_0_v1_0", "pred": " place a slice of ham and a slice of turkey on top of one slice of bread."}]}, "group0123_ex000024_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000024_eng_latn_0_v1_0", "pred": " Tape the letter on the book cover with masking tape."}]}, "group0123_ex000025_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000025_eng_latn_0_v1_0", "pred": " by placing garlic in a plastic bag under cabinet."}]}, "group0123_ex000029_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000029_eng_latn_0_v1_0", "pred": " with small twigs."}]}, "group0123_ex000068_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000068_eng_latn_0_v1_0", "pred": " by pressing the cherries on an upside down bowl."}]}, "group0123_ex000086_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000086_eng_latn_0_v1_0", "pred": " Peel unripened bananas and mash them with a fork."}]}, "group0123_ex000020_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000020_eng_latn_0_v1_0", "pred": " Rub your fingers on stainless steel, like your chopping board."}]}, "group0123_ex000057_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000057_eng_latn_0_v1_0", "pred": " Use a hair dryer to heat up a spot on the PVC pipe and slowly apply pressure on the area you want to bend after it has cooled. Once you've made a bend to the desired angle, you can remove the auger."}]}, "group0123_ex000092_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000092_eng_latn_0_v1_0", "pred": " use a hammer."}]}, "group0123_ex000037_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000037_eng_latn_0_v1_0", "pred": " Place some sausage links in a pot and cover with an inch of water. Simmer for ten minutes or until firm. Allow to cool before serving."}]}, "group0123_ex000044_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000044_eng_latn_0_v1_0", "pred": " Wash jeans on hot cycle using laundry soap and oxyclean, repeat if stain is still there."}]}, "group0123_ex000083_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000083_eng_latn_0_v1_0", "pred": " open one end of the can. Flip the can over and open the bottom. Push the paste out like a Push-Up Pop."}]}, "group0123_ex000008_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000008_eng_latn_0_v1_0", "pred": " Use a pot that holds at least eight quarts."}]}, "group0123_ex000032_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000032_eng_latn_0_v1_0", "pred": " Place the clay sticks onto a baking sheet and bake at 350\u00b0F for twenty minutes."}]}, "group0123_ex000096_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000096_eng_latn_0_v1_0", "pred": " don't try to cut all the way through in one go."}]}, "group0123_ex000084_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000084_eng_latn_0_v1_0", "pred": " Place a wet towel on the counter first."}]}, "group0123_ex000038_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000038_eng_latn_0_v1_0", "pred": " Bolt the bike racks above the cars and other furniture."}]}, "group0123_ex000005_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000005_eng_latn_0_v1_0", "pred": " insert a toothpick into the middle of the bread. If it comes out clean it is done."}]}, "group0123_ex000058_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000058_eng_latn_0_v1_0", "pred": " Lift the page up and push it to the right side."}]}, "group0123_ex000001_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000001_eng_latn_0_v1_0", "pred": " preheat oven to 450 degrees Fahrenheit. Toss the broccoli florets with the olive pits, garlic, salt, and pepper. Spread them out on a baking sheet and then roast, without stirring, until the edges are crispy and the stems are crisp tender."}]}, "group0123_ex000047_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000047_eng_latn_0_v1_0", "pred": " use a solder pen to melt it free and remove by hand."}]}, "group0123_ex000069_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000069_eng_latn_0_v1_0", "pred": " mix together two tablespoons each of red clay, cornstarch, and tapioca pudding and mix well. You can mix in a bit of cocoa powder if you have dark hair."}]}, "group0123_ex000023_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000023_eng_latn_0_v1_0", "pred": " It is best if a skateboard has two wheels to provide balance and traction."}]}, "group0123_ex000098_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000098_eng_latn_0_v1_0", "pred": " spray WD-40 on the the hinges."}]}, "group0123_ex000011_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000011_eng_latn_0_v1_0", "pred": " remove any nails and fill nail holes with spackling."}]}, "group0123_ex000009_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000009_eng_latn_0_v1_0", "pred": " put a piece of plastic wrap under the food on the lower rack."}]}, "group0123_ex000055_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000055_eng_latn_0_v1_0", "pred": " place the cake on the plate then line it with plastic wrap."}]}, "group0123_ex000027_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000027_eng_latn_0_v1_0", "pred": " take a sip of soda."}]}, "group0123_ex000051_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000051_eng_latn_0_v1_0", "pred": " dry your hair with a blow dryer instead of a towel."}]}, "group0123_ex000004_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000004_eng_latn_0_v1_0", "pred": " Clean and gut fish. Sprinkle with salt and pepper. Add about an inch of oil in pan, turn stove on low. Once oil is hot, gently put fish in. Fry each side for 5 minutes until crispy yellow brown."}]}, "group0123_ex000080_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000080_eng_latn_0_v1_0", "pred": " Heat it up on the stove top on high temperature."}]}, "group0123_ex000054_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000054_eng_latn_0_v1_0", "pred": " wait for the edges to brown and the top to bubble. Slide the spatula under the pancake so it's centered. Quickly flip your wrist 90 degrees to one side so the pancake flips over evenly."}]}, "group0123_ex000016_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000016_eng_latn_0_v1_0", "pred": " use rubbing alcohol."}]}, "group0123_ex000066_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000066_eng_latn_0_v1_0", "pred": " boil the beans for a day and then rinse them after cooking."}]}, "group0123_ex000079_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000079_eng_latn_0_v1_0", "pred": " wear dark clothing."}]}, "group0123_ex000074_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000074_eng_latn_0_v1_0", "pred": " wrap the chicken in a damp paper towel and heat in the microwave."}]}, "group0123_ex000090_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000090_eng_latn_0_v1_0", "pred": " Use a bucket with a solid lid."}]}, "group0123_ex000099_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000099_eng_latn_0_v1_0", "pred": " rub a pencil tip on the teeth of the zipper."}]}, "group0123_ex000060_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000060_eng_latn_0_v1_0", "pred": " rub a dryer sheet on your hair."}]}, "group0123_ex000085_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000085_eng_latn_0_v1_0", "pred": " by placing a white sock over the head."}]}, "group0123_ex000043_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000043_eng_latn_0_v1_0", "pred": " mashing up egg whites into a paste."}]}, "group0123_ex000050_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000050_eng_latn_0_v1_0", "pred": " In a large bowl, combine 4 teaspoons half and half light cream, 1 1/2 cups sugar, and 1 tablespoon vanilla. Stir until sugar dissolved. Stir in 2 teaspoons whipping cream. Add 4 cups diced fresh pitted cherries to mixture. Freeze in 4 or 5 quart ice cream freezer according to manufacturers directions. Top with whipped cream and a fresh cherry on top."}]}, "group0123_ex000077_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000077_eng_latn_0_v1_0", "pred": " spray the 9x9 glass pan with cooking spray."}]}, "group0123_ex000078_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000078_eng_latn_0_v1_0", "pred": " put it in the dishwasher."}]}, "group0123_ex000097_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000097_eng_latn_0_v1_0", "pred": " Rinse out blender, and add 1 1/2 cups of hot water and about 1 tablespoon of dish soap. Seal lid then turn blender on high for 1 1/2 minutes."}]}, "group0123_ex000067_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000067_eng_latn_0_v1_0", "pred": " Bolt the bench vise to a wooden work bench."}]}, "group0123_ex000035_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000035_eng_latn_0_v1_0", "pred": " You should check on the BBQ every 45 minutes to make sure it remains between 225-350\u00b0 C."}]}, "group0123_ex000048_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000048_eng_latn_0_v1_0", "pred": " Turn the knob with 'H' on it."}]}, "group0123_ex000040_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000040_eng_latn_0_v1_0", "pred": " sprinkle salt on carpet every nine days and vacuum every third day."}]}, "group0123_ex000003_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000003_eng_latn_0_v1_0", "pred": " put the jar of honey in a microwave until it becomes runny."}]}, "group0123_ex000101_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000101_eng_latn_0_v1_0", "pred": " can be used to make a ghost costume."}]}, "group0123_ex000082_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000082_eng_latn_0_v1_0", "pred": " Cut into it right away."}]}, "group0123_ex000017_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000017_eng_latn_0_v1_0", "pred": " Mix together 1 1/2 cups cornstarch, 1 cup white school glue, 1 tablespoon white vinegar, 1 teaspoon lotion, and 1 tablespoon baby oil. Mix and knead until the desired consistancy."}]}, "group0123_ex000015_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000015_eng_latn_0_v1_0", "pred": " Add 35 drops of blue food coloring and 20 drops of red food coloring to icing. Mix thoroughly before spreading on cake or cupcakes."}]}, "group0123_ex000063_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000063_eng_latn_0_v1_0", "pred": " use a ball-peen hammer."}]}, "group0123_ex000026_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000026_eng_latn_0_v1_0", "pred": " take it out of the oven about 10-15 minutes before you want to roll it out."}]}, "group0123_ex000103_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000103_eng_latn_0_v1_0", "pred": " Use a disposable knife to cut off the pills."}]}, "group0123_ex000049_eng_latn_0_v1": {"predictions": [{"id": "group0123_ex000049_eng_latn_0_v1_0", "pred": " touch the meat and see if it is soft to the touch."}]}}

babylm_eval_repo/strict/results/xpertgpt_fresh/main/zero_shot/causal/global_piqa_parallel/global_piqa_parallel/best_temperature_report.txt

@@ -0,0 +1,110 @@
+TEMPERATURE: 1.00
+
+### UID ACCURACY
+parallel_ex000000_eng_latn: 0.00
+parallel_ex000001_eng_latn: 0.00
+parallel_ex000002_eng_latn: 0.00
+parallel_ex000003_eng_latn: 0.00
+parallel_ex000004_eng_latn: 100.00
+parallel_ex000005_eng_latn: 100.00
+parallel_ex000008_eng_latn: 0.00
+parallel_ex000009_eng_latn: 0.00
+parallel_ex000010_eng_latn: 0.00
+parallel_ex000011_eng_latn: 0.00
+parallel_ex000013_eng_latn: 0.00
+parallel_ex000014_eng_latn: 0.00
+parallel_ex000015_eng_latn: 0.00
+parallel_ex000016_eng_latn: 100.00
+parallel_ex000017_eng_latn: 0.00
+parallel_ex000018_eng_latn: 100.00
+parallel_ex000020_eng_latn: 0.00
+parallel_ex000021_eng_latn: 0.00
+parallel_ex000022_eng_latn: 0.00
+parallel_ex000023_eng_latn: 0.00
+parallel_ex000024_eng_latn: 0.00
+parallel_ex000025_eng_latn: 100.00
+parallel_ex000026_eng_latn: 0.00
+parallel_ex000027_eng_latn: 100.00
+parallel_ex000028_eng_latn: 0.00
+parallel_ex000029_eng_latn: 0.00
+parallel_ex000030_eng_latn: 0.00
+parallel_ex000031_eng_latn: 0.00
+parallel_ex000032_eng_latn: 0.00
+parallel_ex000033_eng_latn: 0.00
+parallel_ex000034_eng_latn: 0.00
+parallel_ex000035_eng_latn: 0.00
+parallel_ex000036_eng_latn: 0.00
+parallel_ex000037_eng_latn: 0.00
+parallel_ex000038_eng_latn: 100.00
+parallel_ex000039_eng_latn: 0.00
+parallel_ex000040_eng_latn: 0.00
+parallel_ex000041_eng_latn: 0.00
+parallel_ex000042_eng_latn: 0.00
+parallel_ex000043_eng_latn: 0.00
+parallel_ex000044_eng_latn: 0.00
+parallel_ex000045_eng_latn: 0.00
+parallel_ex000046_eng_latn: 100.00
+parallel_ex000047_eng_latn: 0.00
+parallel_ex000048_eng_latn: 0.00
+parallel_ex000049_eng_latn: 0.00
+parallel_ex000050_eng_latn: 0.00
+parallel_ex000051_eng_latn: 100.00
+parallel_ex000053_eng_latn: 0.00
+parallel_ex000054_eng_latn: 0.00
+parallel_ex000055_eng_latn: 0.00
+parallel_ex000056_eng_latn: 0.00
+parallel_ex000057_eng_latn: 0.00
+parallel_ex000058_eng_latn: 100.00
+parallel_ex000059_eng_latn: 0.00
+parallel_ex000061_eng_latn: 0.00
+parallel_ex000062_eng_latn: 0.00
+parallel_ex000063_eng_latn: 0.00
+parallel_ex000064_eng_latn: 0.00
+parallel_ex000065_eng_latn: 0.00
+parallel_ex000066_eng_latn: 0.00
+parallel_ex000067_eng_latn: 0.00
+parallel_ex000068_eng_latn: 100.00
+parallel_ex000069_eng_latn: 0.00
+parallel_ex000070_eng_latn: 0.00
+parallel_ex000071_eng_latn: 0.00
+parallel_ex000072_eng_latn: 100.00
+parallel_ex000073_eng_latn: 0.00
+parallel_ex000074_eng_latn: 100.00
+parallel_ex000075_eng_latn: 0.00
+parallel_ex000076_eng_latn: 100.00
+parallel_ex000077_eng_latn: 100.00
+parallel_ex000078_eng_latn: 0.00
+parallel_ex000079_eng_latn: 0.00
+parallel_ex000080_eng_latn: 0.00
+parallel_ex000081_eng_latn: 100.00
+parallel_ex000082_eng_latn: 0.00
+parallel_ex000083_eng_latn: 0.00
+parallel_ex000084_eng_latn: 0.00
+parallel_ex000085_eng_latn: 0.00
+parallel_ex000086_eng_latn: 100.00
+parallel_ex000087_eng_latn: 0.00
+parallel_ex000088_eng_latn: 0.00
+parallel_ex000089_eng_latn: 0.00
+parallel_ex000090_eng_latn: 0.00
+parallel_ex000091_eng_latn: 0.00
+parallel_ex000092_eng_latn: 0.00
+parallel_ex000093_eng_latn: 0.00
+parallel_ex000094_eng_latn: 0.00
+parallel_ex000095_eng_latn: 0.00
+parallel_ex000096_eng_latn: 0.00
+parallel_ex000097_eng_latn: 0.00
+parallel_ex000098_eng_latn: 100.00
+parallel_ex000099_eng_latn: 0.00
+parallel_ex000100_eng_latn: 0.00
+parallel_ex000101_eng_latn: 0.00
+parallel_ex000102_eng_latn: 0.00
+parallel_ex000103_eng_latn: 0.00
+parallel_ex000104_eng_latn: 0.00
+parallel_ex000105_eng_latn: 100.00
+parallel_ex000106_eng_latn: 100.00
+parallel_ex000107_eng_latn: 0.00
+parallel_ex000108_eng_latn: 0.00
+
+### AVERAGE ACCURACY
+19.42
+

babylm_eval_repo/strict/results/xpertgpt_fresh/main/zero_shot/causal/global_piqa_parallel/global_piqa_parallel/predictions.json

@@ -0,0 +1 @@
+{"parallel_ex000000_eng_latn": {"predictions": [{"id": "parallel_ex000000_eng_latn_0", "pred": " The amount of air in the bag increases"}]}, "parallel_ex000001_eng_latn": {"predictions": [{"id": "parallel_ex000001_eng_latn_0", "pred": " The water goes into the cup"}]}, "parallel_ex000002_eng_latn": {"predictions": [{"id": "parallel_ex000002_eng_latn_0", "pred": " The ball sticks to the window"}]}, "parallel_ex000003_eng_latn": {"predictions": [{"id": "parallel_ex000003_eng_latn_0", "pred": " The ball sticks to the window"}]}, "parallel_ex000004_eng_latn": {"predictions": [{"id": "parallel_ex000004_eng_latn_0", "pred": " You see the light from the lightbulb through the glass"}]}, "parallel_ex000005_eng_latn": {"predictions": [{"id": "parallel_ex000005_eng_latn_0", "pred": " The center of the fabric gets closer to the ground"}]}, "parallel_ex000008_eng_latn": {"predictions": [{"id": "parallel_ex000008_eng_latn_0", "pred": " October"}]}, "parallel_ex000009_eng_latn": {"predictions": [{"id": "parallel_ex000009_eng_latn_0", "pred": " Friday"}]}, "parallel_ex000010_eng_latn": {"predictions": [{"id": "parallel_ex000010_eng_latn_0", "pred": " 1940"}]}, "parallel_ex000011_eng_latn": {"predictions": [{"id": "parallel_ex000011_eng_latn_0", "pred": " 1985"}]}, "parallel_ex000013_eng_latn": {"predictions": [{"id": "parallel_ex000013_eng_latn_0", "pred": " 9:30 AM"}]}, "parallel_ex000014_eng_latn": {"predictions": [{"id": "parallel_ex000014_eng_latn_0", "pred": " North"}]}, "parallel_ex000015_eng_latn": {"predictions": [{"id": "parallel_ex000015_eng_latn_0", "pred": " Late morning"}]}, "parallel_ex000016_eng_latn": {"predictions": [{"id": "parallel_ex000016_eng_latn_0", "pred": " You pull it"}]}, "parallel_ex000017_eng_latn": {"predictions": [{"id": "parallel_ex000017_eng_latn_0", "pred": " A stick"}]}, "parallel_ex000018_eng_latn": {"predictions": [{"id": "parallel_ex000018_eng_latn_0", "pred": " Watching a movie"}]}, "parallel_ex000020_eng_latn": {"predictions": [{"id": "parallel_ex000020_eng_latn_0", "pred": " A plate"}]}, "parallel_ex000021_eng_latn": {"predictions": [{"id": "parallel_ex000021_eng_latn_0", "pred": " Lemons"}]}, "parallel_ex000022_eng_latn": {"predictions": [{"id": "parallel_ex000022_eng_latn_0", "pred": " Marbles"}]}, "parallel_ex000023_eng_latn": {"predictions": [{"id": "parallel_ex000023_eng_latn_0", "pred": " A wooden box"}]}, "parallel_ex000024_eng_latn": {"predictions": [{"id": "parallel_ex000024_eng_latn_0", "pred": " A small stone"}]}, "parallel_ex000025_eng_latn": {"predictions": [{"id": "parallel_ex000025_eng_latn_0", "pred": " A piece of rope"}]}, "parallel_ex000026_eng_latn": {"predictions": [{"id": "parallel_ex000026_eng_latn_0", "pred": " Under the table"}]}, "parallel_ex000027_eng_latn": {"predictions": [{"id": "parallel_ex000027_eng_latn_0", "pred": " Inside an empty suitcase"}]}, "parallel_ex000028_eng_latn": {"predictions": [{"id": "parallel_ex000028_eng_latn_0", "pred": " Your fingers"}]}, "parallel_ex000029_eng_latn": {"predictions": [{"id": "parallel_ex000029_eng_latn_0", "pred": " Ten pairs of glasses"}]}, "parallel_ex000030_eng_latn": {"predictions": [{"id": "parallel_ex000030_eng_latn_0", "pred": " It depends on how ripe the spinach is"}]}, "parallel_ex000031_eng_latn": {"predictions": [{"id": "parallel_ex000031_eng_latn_0", "pred": " It depends on the color of the water"}]}, "parallel_ex000032_eng_latn": {"predictions": [{"id": "parallel_ex000032_eng_latn_0", "pred": " You"}]}, "parallel_ex000033_eng_latn": {"predictions": [{"id": "parallel_ex000033_eng_latn_0", "pred": " Keep the comb still"}]}, "parallel_ex000034_eng_latn": {"predictions": [{"id": "parallel_ex000034_eng_latn_0", "pred": " Your head"}]}, "parallel_ex000035_eng_latn": {"predictions": [{"id": "parallel_ex000035_eng_latn_0", "pred": " It will be crushed"}]}, "parallel_ex000036_eng_latn": {"predictions": [{"id": "parallel_ex000036_eng_latn_0", "pred": " It will break"}]}, "parallel_ex000037_eng_latn": {"predictions": [{"id": "parallel_ex000037_eng_latn_0", "pred": " Green"}]}, "parallel_ex000038_eng_latn": {"predictions": [{"id": "parallel_ex000038_eng_latn_0", "pred": " Metal"}]}, "parallel_ex000039_eng_latn": {"predictions": [{"id": "parallel_ex000039_eng_latn_0", "pred": " Chew without swallowing"}]}, "parallel_ex000040_eng_latn": {"predictions": [{"id": "parallel_ex000040_eng_latn_0", "pred": " Large and thin"}]}, "parallel_ex000041_eng_latn": {"predictions": [{"id": "parallel_ex000041_eng_latn_0", "pred": " The man and the woman are on the same floor"}]}, "parallel_ex000042_eng_latn": {"predictions": [{"id": "parallel_ex000042_eng_latn_0", "pred": " The man and the woman are on the same floor"}]}, "parallel_ex000043_eng_latn": {"predictions": [{"id": "parallel_ex000043_eng_latn_0", "pred": " The paper is on top of the book"}]}, "parallel_ex000044_eng_latn": {"predictions": [{"id": "parallel_ex000044_eng_latn_0", "pred": " A half cup"}]}, "parallel_ex000045_eng_latn": {"predictions": [{"id": "parallel_ex000045_eng_latn_0", "pred": " Your friend is closer to you"}]}, "parallel_ex000046_eng_latn": {"predictions": [{"id": "parallel_ex000046_eng_latn_0", "pred": " The rock is no longer visible"}]}, "parallel_ex000047_eng_latn": {"predictions": [{"id": "parallel_ex000047_eng_latn_0", "pred": " The tunnel is no longer visible from the car"}]}, "parallel_ex000048_eng_latn": {"predictions": [{"id": "parallel_ex000048_eng_latn_0", "pred": " From the side of the shelf"}]}, "parallel_ex000049_eng_latn": {"predictions": [{"id": "parallel_ex000049_eng_latn_0", "pred": " The can is not visible in front of the flour"}]}, "parallel_ex000050_eng_latn": {"predictions": [{"id": "parallel_ex000050_eng_latn_0", "pred": " The coin is dented"}]}, "parallel_ex000051_eng_latn": {"predictions": [{"id": "parallel_ex000051_eng_latn_0", "pred": " The cup is crushed"}]}, "parallel_ex000053_eng_latn": {"predictions": [{"id": "parallel_ex000053_eng_latn_0", "pred": " The library doesn't have any books"}]}, "parallel_ex000054_eng_latn": {"predictions": [{"id": "parallel_ex000054_eng_latn_0", "pred": " There are no students in classroom B"}]}, "parallel_ex000055_eng_latn": {"predictions": [{"id": "parallel_ex000055_eng_latn_0", "pred": " The theatre has more chairs than the school"}]}, "parallel_ex000056_eng_latn": {"predictions": [{"id": "parallel_ex000056_eng_latn_0", "pred": " All of the cups are in the cupboard"}]}, "parallel_ex000057_eng_latn": {"predictions": [{"id": "parallel_ex000057_eng_latn_0", "pred": " 16cm x 20cm"}]}, "parallel_ex000058_eng_latn": {"predictions": [{"id": "parallel_ex000058_eng_latn_0", "pred": " 1cm"}]}, "parallel_ex000059_eng_latn": {"predictions": [{"id": "parallel_ex000059_eng_latn_0", "pred": " A knife"}]}, "parallel_ex000061_eng_latn": {"predictions": [{"id": "parallel_ex000061_eng_latn_0", "pred": " They should not twist the towel"}]}, "parallel_ex000062_eng_latn": {"predictions": [{"id": "parallel_ex000062_eng_latn_0", "pred": " Your left shoulder"}]}, "parallel_ex000063_eng_latn": {"predictions": [{"id": "parallel_ex000063_eng_latn_0", "pred": " A foam cup"}]}, "parallel_ex000064_eng_latn": {"predictions": [{"id": "parallel_ex000064_eng_latn_0", "pred": " The side of the house that faces the sun"}]}, "parallel_ex000065_eng_latn": {"predictions": [{"id": "parallel_ex000065_eng_latn_0", "pred": " Without oven mitts, grab it from both sides and carefully lift it out of the oven"}]}, "parallel_ex000066_eng_latn": {"predictions": [{"id": "parallel_ex000066_eng_latn_0", "pred": " A pocket knife"}]}, "parallel_ex000067_eng_latn": {"predictions": [{"id": "parallel_ex000067_eng_latn_0", "pred": " Scissors"}]}, "parallel_ex000068_eng_latn": {"predictions": [{"id": "parallel_ex000068_eng_latn_0", "pred": " Pants, socks, then shoes"}]}, "parallel_ex000069_eng_latn": {"predictions": [{"id": "parallel_ex000069_eng_latn_0", "pred": " Heavy wool shirts"}]}, "parallel_ex000070_eng_latn": {"predictions": [{"id": "parallel_ex000070_eng_latn_0", "pred": " Drink all your water as soon as you start the hike"}]}, "parallel_ex000071_eng_latn": {"predictions": [{"id": "parallel_ex000071_eng_latn_0", "pred": " The wick should be as long as possible"}]}, "parallel_ex000072_eng_latn": {"predictions": [{"id": "parallel_ex000072_eng_latn_0", "pred": " They will make the same amount of water overflow"}]}, "parallel_ex000073_eng_latn": {"predictions": [{"id": "parallel_ex000073_eng_latn_0", "pred": " A water-soaked cotton ball"}]}, "parallel_ex000074_eng_latn": {"predictions": [{"id": "parallel_ex000074_eng_latn_0", "pred": " Wrap the cup with a plastic bag"}]}, "parallel_ex000075_eng_latn": {"predictions": [{"id": "parallel_ex000075_eng_latn_0", "pred": " 10 liters"}]}, "parallel_ex000076_eng_latn": {"predictions": [{"id": "parallel_ex000076_eng_latn_0", "pred": " Do not cover it"}]}, "parallel_ex000077_eng_latn": {"predictions": [{"id": "parallel_ex000077_eng_latn_0", "pred": " A stick"}]}, "parallel_ex000078_eng_latn": {"predictions": [{"id": "parallel_ex000078_eng_latn_0", "pred": " In the sun"}]}, "parallel_ex000079_eng_latn": {"predictions": [{"id": "parallel_ex000079_eng_latn_0", "pred": " The car has lost a tire"}]}, "parallel_ex000080_eng_latn": {"predictions": [{"id": "parallel_ex000080_eng_latn_0", "pred": " They will freeze at the same rate"}]}, "parallel_ex000081_eng_latn": {"predictions": [{"id": "parallel_ex000081_eng_latn_0", "pred": " Flatten out the towel and do not fold it"}]}, "parallel_ex000082_eng_latn": {"predictions": [{"id": "parallel_ex000082_eng_latn_0", "pred": " A high rattling sound"}]}, "parallel_ex000083_eng_latn": {"predictions": [{"id": "parallel_ex000083_eng_latn_0", "pred": " More than 20cm"}]}, "parallel_ex000084_eng_latn": {"predictions": [{"id": "parallel_ex000084_eng_latn_0", "pred": " Yes"}]}, "parallel_ex000085_eng_latn": {"predictions": [{"id": "parallel_ex000085_eng_latn_0", "pred": " Make animal shapes with your hand, then place your hand on the wall"}]}, "parallel_ex000086_eng_latn": {"predictions": [{"id": "parallel_ex000086_eng_latn_0", "pred": " The glass would break"}]}, "parallel_ex000087_eng_latn": {"predictions": [{"id": "parallel_ex000087_eng_latn_0", "pred": " Sugar and water, flour and salt"}]}, "parallel_ex000088_eng_latn": {"predictions": [{"id": "parallel_ex000088_eng_latn_0", "pred": " Two glasses of 0.2L and one glass of 0.5L"}]}, "parallel_ex000089_eng_latn": {"predictions": [{"id": "parallel_ex000089_eng_latn_0", "pred": " Put the egg in the cup, then wrap the cup with the paper"}]}, "parallel_ex000090_eng_latn": {"predictions": [{"id": "parallel_ex000090_eng_latn_0", "pred": " The water will not move"}]}, "parallel_ex000091_eng_latn": {"predictions": [{"id": "parallel_ex000091_eng_latn_0", "pred": " Blue and white"}]}, "parallel_ex000092_eng_latn": {"predictions": [{"id": "parallel_ex000092_eng_latn_0", "pred": " One drop of yellow and one drop of blue"}]}, "parallel_ex000093_eng_latn": {"predictions": [{"id": "parallel_ex000093_eng_latn_0", "pred": " Fill a bag with water"}]}, "parallel_ex000094_eng_latn": {"predictions": [{"id": "parallel_ex000094_eng_latn_0", "pred": " Folding the laundry"}]}, "parallel_ex000095_eng_latn": {"predictions": [{"id": "parallel_ex000095_eng_latn_0", "pred": " In front of you"}]}, "parallel_ex000096_eng_latn": {"predictions": [{"id": "parallel_ex000096_eng_latn_0", "pred": " You and your friend are walking towards each other"}]}, "parallel_ex000097_eng_latn": {"predictions": [{"id": "parallel_ex000097_eng_latn_0", "pred": " The right car would swerve right, and the left car would not swerve"}]}, "parallel_ex000098_eng_latn": {"predictions": [{"id": "parallel_ex000098_eng_latn_0", "pred": " 240km"}]}, "parallel_ex000099_eng_latn": {"predictions": [{"id": "parallel_ex000099_eng_latn_0", "pred": " 9:30 AM"}]}, "parallel_ex000100_eng_latn": {"predictions": [{"id": "parallel_ex000100_eng_latn_0", "pred": " Friday"}]}, "parallel_ex000101_eng_latn": {"predictions": [{"id": "parallel_ex000101_eng_latn_0", "pred": " 9:00 AM"}]}, "parallel_ex000102_eng_latn": {"predictions": [{"id": "parallel_ex000102_eng_latn_0", "pred": " Cooking oil"}]}, "parallel_ex000103_eng_latn": {"predictions": [{"id": "parallel_ex000103_eng_latn_0", "pred": " Leave the garlic in the sun"}]}, "parallel_ex000104_eng_latn": {"predictions": [{"id": "parallel_ex000104_eng_latn_0", "pred": " Small rocks"}]}, "parallel_ex000105_eng_latn": {"predictions": [{"id": "parallel_ex000105_eng_latn_0", "pred": " The meat is soft"}]}, "parallel_ex000106_eng_latn": {"predictions": [{"id": "parallel_ex000106_eng_latn_0", "pred": " A spoon"}]}, "parallel_ex000107_eng_latn": {"predictions": [{"id": "parallel_ex000107_eng_latn_0", "pred": " A heavy coat"}]}, "parallel_ex000108_eng_latn": {"predictions": [{"id": "parallel_ex000108_eng_latn_0", "pred": " A drinking cup"}]}}

babylm_eval_repo/strict/results/xpertgpt_fresh/main/zero_shot/causal/reading/correlations.txt

@@ -0,0 +1,11 @@
+RTfirstfix	0.5693
+RTfirstpass	0.5957
+RTgopast	0.4633
+RTrightbound	0.5995
+self_paced_reading_time	0.2430
+ELAN	0.1066
+LAN	-0.0515
+N400	-0.1010
+P600	0.2130
+EPNP	0.0968
+PNP	0.2318

babylm_eval_repo/strict/results/xpertgpt_fresh/main/zero_shot/causal/reading/predictive_power.jsonl

@@ -0,0 +1,11 @@
+{"Predicted variable": "RTfirstfix", "Coefficient": 1.3490699037631342, "Number of standard deviations": 4.195750408501817, "P-value": 2.861052285519171e-05, "R2": 0.6477540607490784, "Change in R2 from baseline": 0.003673608269121109, "AIC": 16389.11553716685, "Change in AIC from baseline": -15.596156912710285}
+{"Predicted variable": "RTfirstpass", "Coefficient": 2.4611804773850223, "Number of standard deviations": 5.501074574866903, "P-value": 4.3542625538347874e-08, "R2": 0.650588875740541, "Change in R2 from baseline": 0.0062641096355765535, "AIC": 17509.67787915891, "Change in AIC from baseline": -28.135909857664956}
+{"Predicted variable": "RTgopast", "Coefficient": 12.305662283836849, "Number of standard deviations": 9.33835253209628, "P-value": 2.9808239336644636e-20, "R2": 0.5617358589628113, "Change in R2 from baseline": 0.022641439006945485, "AIC": 21173.824833937466, "Change in AIC from baseline": -83.42992180746296}
+{"Predicted variable": "RTrightbound", "Coefficient": 4.37062955152205, "Number of standard deviations": 7.810200157453945, "P-value": 9.953705102336702e-15, "R2": 0.6319697560710787, "Change in R2 from baseline": 0.013299502492940984, "AIC": 18268.735621979176, "Change in AIC from baseline": -58.206910505876294}
+{"Predicted variable": "self_paced_reading_time", "Coefficient": 1.7185264038913184, "Number of standard deviations": 6.837599172601403, "P-value": 1.1217426053972441e-11, "R2": 0.09324041708333053, "Change in R2 from baseline": 0.02511464180981937, "AIC": 15553.622928507544, "Change in AIC from baseline": -44.33557993874456}
+{"Predicted variable": "ELAN", "Coefficient": -0.026772643700025297, "Number of standard deviations": -2.3100926986452808, "P-value": 0.021003127045723314, "R2": 0.06658383148523661, "Change in R2 from baseline": 0.0029509489199446293, "AIC": 5117.480838655834, "Change in AIC from baseline": -3.353362136870601}
+{"Predicted variable": "LAN", "Coefficient": -0.005468477371309654, "Number of standard deviations": -0.40851411308016833, "P-value": 0.6829480877546681, "R2": 0.12204982479130544, "Change in R2 from baseline": 8.679836753788805e-05, "AIC": 5606.390417937283, "Change in AIC from baseline": 1.8323335890399903}
+{"Predicted variable": "N400", "Coefficient": -0.036789487463063514, "Number of standard deviations": -3.2812512939025096, "P-value": 0.0010545563469490558, "R2": 0.16739385815162078, "Change in R2 from baseline": 0.005310631313760639, "AIC": 5005.186440596301, "Change in AIC from baseline": -8.783283448517977}
+{"Predicted variable": "P600", "Coefficient": 0.028012199337309463, "Number of standard deviations": 2.8988489951015777, "P-value": 0.0037937324510707417, "R2": 0.20731173471252906, "Change in R2 from baseline": 0.003946218904258503, "AIC": 4500.928498302437, "Change in AIC from baseline": -6.422205014940118}
+{"Predicted variable": "EPNP", "Coefficient": 0.03623777396385926, "Number of standard deviations": 3.8573259094568546, "P-value": 0.00011894202763811692, "R2": 0.11988027893731101, "Change in R2 from baseline": 0.007757860732487076, "AIC": 4405.279303931775, "Change in AIC from baseline": -12.883977736527413}
+{"Predicted variable": "PNP", "Coefficient": 0.059273137039820734, "Number of standard deviations": 5.775466017324852, "P-value": 9.1171100402867e-09, "R2": 0.1833428383886203, "Change in R2 from baseline": 0.01613769110490293, "AIC": 4705.163896945838, "Change in AIC from baseline": -31.187261221406516}

babylm_eval_repo/strict/results/xpertgpt_fresh/main/zero_shot/causal/reading/predictive_power_spillover.jsonl

@@ -0,0 +1,11 @@
+{"Predicted variable": "RTfirstfix", "Coefficient": 1.6217153954117451, "Number of standard deviations": 4.841745887039933, "P-value": 1.4219074334063838e-06, "R2": 0.7071365504941127, "Change in R2 from baseline": 0.004641955620255711, "AIC": 14298.852913604147, "Change in AIC from baseline": -21.52599961041051}
+{"Predicted variable": "RTfirstpass", "Coefficient": 2.8014033494515846, "Number of standard deviations": 6.149090535941275, "P-value": 1.0007288324286125e-09, "R2": 0.7067820953051911, "Change in R2 from baseline": 0.007496250431207052, "AIC": 15219.213651155645, "Change in AIC from baseline": -35.76521573182072}
+{"Predicted variable": "RTgopast", "Coefficient": 13.4660921862215, "Number of standard deviations": 9.004183847817785, "P-value": 6.518290395675876e-19, "R2": 0.586109834183481, "Change in R2 from baseline": 0.02268849252097327, "AIC": 18775.711889611594, "Change in AIC from baseline": -77.83841244154974}
+{"Predicted variable": "RTrightbound", "Coefficient": 4.711070672096454, "Number of standard deviations": 7.912726576139806, "P-value": 4.8969133890847694e-15, "R2": 0.6744689608176773, "Change in R2 from baseline": 0.013780867227876747, "AIC": 16019.952083886219, "Change in AIC from baseline": -60.02707226482562}
+{"Predicted variable": "self_paced_reading_time", "Coefficient": 1.3743440893458203, "Number of standard deviations": 5.2317550238035, "P-value": 1.9201083048415453e-07, "R2": 0.14748370377557674, "Change in R2 from baseline": 0.015777177643346074, "AIC": 13571.85885415219, "Change in AIC from baseline": -25.432804659596513}
+{"Predicted variable": "ELAN", "Coefficient": -0.022321853267920702, "Number of standard deviations": -1.7807583439800752, "P-value": 0.07515714344715244, "R2": 0.08033606732135334, "Change in R2 from baseline": 0.0019718367505782153, "AIC": 4468.823848831604, "Change in AIC from baseline": -1.2041159873515426}
+{"Predicted variable": "LAN", "Coefficient": 0.007777264243250142, "Number of standard deviations": 0.5232766881289969, "P-value": 0.6008601084277922, "R2": 0.12876126826798107, "Change in R2 from baseline": 0.0001612990371812062, "AIC": 4978.429351711803, "Change in AIC from baseline": 1.723059803896831}
+{"Predicted variable": "N400", "Coefficient": -0.0337406160300188, "Number of standard deviations": -2.717324927039772, "P-value": 0.0066579780541857125, "R2": 0.1865806933643791, "Change in R2 from baseline": 0.004060966881984318, "AIC": 4440.482323832297, "Change in AIC from baseline": -5.45014479765905}
+{"Predicted variable": "P600", "Coefficient": 0.038339925099556346, "Number of standard deviations": 3.559916890335394, "P-value": 0.00038271288738986304, "R2": 0.22402150422229639, "Change in R2 from baseline": 0.006649074977216229, "AIC": 4014.721212766805, "Change in AIC from baseline": -10.76406748618183}
+{"Predicted variable": "EPNP", "Coefficient": 0.05247397576153977, "Number of standard deviations": 5.0011567707408995, "P-value": 6.379359738302813e-07, "R2": 0.11019784528001031, "Change in R2 from baseline": 0.015047564590695006, "AIC": 3936.610631317539, "Change in AIC from baseline": -23.08752199507444}
+{"Predicted variable": "PNP", "Coefficient": 0.0747711475299754, "Number of standard deviations": 6.496567833073147, "P-value": 1.119714283869015e-10, "R2": 0.19719952360571513, "Change in R2 from baseline": 0.022909067002225836, "AIC": 4213.400442108958, "Change in AIC from baseline": -40.09274101206756}

babylm_eval_repo/strict/results/xpertgpt_fresh/main/zero_shot/causal/reading/report.txt

@@ -0,0 +1,2 @@
+EYE TRACKING SCORE: 2.80
+SELF-PACED READING SCORE: 1.82

babylm_eval_repo/strict/scripts/aoa_bos_curves.py

@@ -0,0 +1,203 @@
+#!/usr/bin/env python
+"""BOS-only AoA: extract P(word | <bos>) across checkpoints (single-BOS, fixed path)
+and plot learning-curve small-multiples by category (averaging + k*ln(V) ceiling +
+robust fit — the final scorer mechanism).
+
+Run from babylm-eval-hidden/strict, or set BABYLM_STRICT to that path.
+Point HF_HOME at scratch if disk is tight; PURGE_PER_CKPT keeps only one revision at a time.
+"""
+import os
+import sys
+import json
+import shutil
+from collections import Counter
+from pathlib import Path
+
+import numpy as np
+
+# ---------------- config ----------------
+MODEL = "BabyLM-community/BabyLM-2026-Baseline-GPT2-Strict"
+BACKEND = "causal"           # causal / mlm / mntp / enc_dec_mask / enc_dec_prefix
+TRACK = "non-strict-small"   # "strict-small" (19 ckpts) or "non-strict-small" (28)
+MIN_CONTEXT = 20             # word set: same filter as the context run
+PURGE_PER_CKPT = True        # delete each HF revision after use (low disk)
+EVAL_DATASET = "BabyLM-community/BabyLM-2026-Strict-Evals"  # source of the AoA/CDI data
+REPO_STRICT = Path(os.environ.get("BABYLM_STRICT", ".")).resolve()
+OUT_DIR = REPO_STRICT / "results" / (Path(MODEL).stem + "_BOS")
+# ----------------------------------------
+
+sys.path.insert(0, str(REPO_STRICT))
+from evaluation_pipeline.AoA_word.eval_util import StepConfig, load_eval  # noqa: E402
+from evaluation_pipeline.AoA_word.evaluation_functions import StepSurprisalExtractor  # noqa: E402
+from evaluation_pipeline.utils import AoAEvaluator  # noqa: E402
+
+WORD_PATH = REPO_STRICT / "evaluation_data/full_eval/aoa/cdi_childes.json"
+CDI = REPO_STRICT / "evaluation_data/full_eval/aoa/cdi_human.csv"
+SURPRISAL = OUT_DIR / "surprisal.json"
+PNG = OUT_DIR / "aoa_bos_curves.png"
+
+
+def ensure_data():
+    # Fetch the AoA/CDI eval files from HF if they aren't present locally.
+    if WORD_PATH.exists() and CDI.exists():
+        return
+    from huggingface_hub import snapshot_download
+    print(f"[data] downloading AoA eval data from {EVAL_DATASET}", flush=True)
+    snapshot_download(
+        repo_id=EVAL_DATASET,
+        repo_type="dataset",
+        local_dir=str(REPO_STRICT),
+        allow_patterns=["evaluation_data/full_eval/aoa/*"],
+    )
+
+
+def sigmoid(x, a, b, c, d):
+    return a / (1 + np.exp(-b * (x - c))) + d
+
+
+def extract():
+    import torch
+    hub = Path(os.environ["HF_HOME"]) / "hub" if os.environ.get("HF_HOME") else None
+
+    def purge():
+        if PURGE_PER_CKPT and hub and hub.exists():
+            shutil.rmtree(hub, ignore_errors=True)
+
+    dev = "cuda" if torch.cuda.is_available() else "cpu"
+    words, _ = load_eval(WORD_PATH, MIN_CONTEXT, debug=False)
+    cfg = StepConfig(resume=False, track=TRACK, file_path=None, debug=False)
+    ex = StepSurprisalExtractor(config=cfg, model_name=MODEL, backend=BACKEND, device=dev)
+    print(f"[extract] {MODEL}  words={len(words)}  steps={len(cfg.steps)}  device={dev}", flush=True)
+    rows = []
+    for step, wc in zip(cfg.steps, cfg.word_counts):
+        purge()
+        print(f"[extract] {step}", flush=True)
+        model = ex.load_model_for_step(step)
+        proc, tok = ex.load_tokenizer_for_step(step)
+        for w in words:
+            s = ex.compute_surprisal(model, proc, tok, "", w, use_bos_only=True)
+            rows.append({"step": step, "word_count": wc, "target_word": w,
+                         "surprisal": s, "use_bos_only": True})
+    purge()
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    json.dump({"results": rows}, SURPRISAL.open("w"))
+    print(f"[extract] DONE rows={len(rows)} -> {SURPRISAL}", flush=True)
+
+
+def robust_fit(ls, neg):
+    from scipy.optimize import curve_fit
+    rng = neg.max() - neg.min()
+    p0 = [rng, 1.0, ls.mean(), neg.min()]
+    lb = [0.0, 0.0, ls.min() - 1, neg.min() - 2 * rng - 1]
+    ub = [10 * rng + 1, 100.0, ls.max() + 1, neg.max() + 1]
+    return curve_fit(sigmoid, ls, neg, p0=p0, bounds=(lb, ub), maxfev=20000)
+
+
+def classify(us, means, k, V):
+    chance = k * np.log(V)
+    thr = chance - 0.5 * (chance - means.min())
+    neg = -means
+    ls = np.log10(us + 1)
+    try:
+        popt, _ = robust_fit(ls, neg)
+    except Exception:
+        return None, "fit-fail", None, thr
+    a, b, c, d = popt
+    nt = -thr
+    if b <= 1e-6 or a <= 1e-6:
+        return None, "bad-shape", popt, thr
+    if nt <= d:
+        return None, "already-acq", popt, thr
+    if nt >= a + d:
+        return None, "never-reached", popt, thr
+    la = c - np.log((a / (nt - d)) - 1) / b
+    if 10 ** la - 1 < us[0] or 10 ** la - 1 > us[-1]:
+        return None, "outside-range", popt, thr
+    return la, "ok", popt, thr
+
+
+def plot():
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    from transformers import AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(MODEL)
+    V = tok.vocab_size
+    pref = len(tok("The", add_special_tokens=False)["input_ids"])
+
+    def klen(w):
+        return max(1, len(tok("The " + w, add_special_tokens=False)["input_ids"]) - pref)
+
+    ev = AoAEvaluator(CDI)
+    wd = {}
+    for r in json.load(SURPRISAL.open())["results"]:
+        st = ev.extract_step_number(r["step"])
+        wd.setdefault(r["target_word"], {"steps": [], "s": []})
+        wd[r["target_word"]]["steps"].append(st)
+        wd[r["target_word"]]["s"].append(float(r["surprisal"]))
+
+    info = {}
+    for w, d in wd.items():
+        us = np.array(sorted(set(d["steps"])), float)
+        means = np.array([np.mean([x for s2, x in zip(d["steps"], d["s"]) if s2 == u]) for u in us])
+        aoa, cat, popt, thr = classify(us, means, klen(w), V)
+        info[w] = dict(us=us, means=means, cat=cat, popt=popt, aoa=aoa, thr=thr,
+                       chance=klen(w) * np.log(V), minm=float(means.min()), k=klen(w))
+    print("category counts:", dict(Counter(i["cat"] for i in info.values()).most_common()))
+
+    C_MEAN, C_RAW, C_FIT, C_THR, C_ANCH, C_CROSS = "#111", "#CCC", "#0072B2", "#009E73", "#999", "#D55E00"
+    counts = Counter(i["cat"] for i in info.values())
+    cats = ["ok"] + [c for c, _ in counts.most_common() if c != "ok"][:4]
+    ncol = 3
+    fig, axes = plt.subplots(len(cats), ncol, figsize=(11, 2.6 * len(cats)),
+                             constrained_layout=True, squeeze=False)
+    for ri, cat in enumerate(cats):
+        words = sorted([w for w, i in info.items() if i["cat"] == cat])[:ncol]
+        for ci in range(ncol):
+            ax = axes[ri][ci]
+            if ci >= len(words):
+                ax.axis("off")
+                continue
+            w = words[ci]; I = info[w]; d = wd[w]
+            st = np.array(d["steps"], float); s = np.array(d["s"], float)
+            ax.scatter(np.log10(st), s, s=5, c=C_RAW, alpha=0.35, zorder=1, linewidths=0)
+            ax.plot(np.log10(I["us"]), I["means"], "-o", color=C_MEAN, ms=3, lw=1.5, zorder=4)
+            if I["popt"] is not None:
+                xs = np.linspace(np.log10(I["us"].min() + 1), np.log10(I["us"].max() + 1), 200)
+                ax.plot(xs, -sigmoid(xs, *I["popt"]), color=C_FIT, lw=1.6, zorder=3)
+            ax.axhline(I["chance"], color=C_ANCH, ls=":", lw=1.0)
+            ax.axhline(I["minm"], color=C_ANCH, ls=":", lw=1.0)
+            ax.axhline(I["thr"], color=C_THR, lw=1.6)
+            if I["aoa"] is not None:
+                ax.axvline(I["aoa"], color=C_CROSS, ls="--", lw=1.2)
+            ax.set_title(f"{w}  (k={I['k']})", fontsize=8.5)
+            ax.grid(True, color="#F0F0F0", lw=0.6); ax.set_axisbelow(True)
+            for sp in ("top", "right"):
+                ax.spines[sp].set_visible(False)
+            ax.tick_params(labelsize=7)
+            if ci == 0:
+                ax.set_ylabel(f"{cat}\nsurprisal (nats)", fontsize=8)
+    for ci in range(ncol):
+        axes[-1][ci].set_xlabel("log10(training words)", fontsize=8)
+    handles = [
+        plt.Line2D([], [], color=C_MEAN, marker="o", ms=3, lw=1.5, label="mean surprisal / ckpt"),
+        plt.Line2D([], [], color=C_FIT, lw=1.6, label="robust sigmoid fit"),
+        plt.Line2D([], [], color=C_THR, lw=1.6, label="threshold = midpoint"),
+        plt.Line2D([], [], color=C_ANCH, ls=":", lw=1.0, label="anchors: k·ln(V) & min"),
+        plt.Line2D([], [], color=C_CROSS, ls="--", lw=1.2, label="crossing (AoA)"),
+    ]
+    fig.legend(handles=handles, loc="upper center", ncol=5, fontsize=8,
+               bbox_to_anchor=(0.5, 1.02), frameon=False)
+    fig.suptitle(f"BOS-only AoA curves  —  {Path(MODEL).stem}", y=1.04, fontsize=11)
+    fig.savefig(PNG, dpi=130, bbox_inches="tight")
+    print(f"[plot] -> {PNG}")
+
+
+if __name__ == "__main__":
+    ensure_data()
+    if not SURPRISAL.exists():
+        extract()
+    else:
+        print(f"[extract] {SURPRISAL} exists, skipping extraction")
+    plot()

babylm_eval_repo/strict/scripts/collate_preds.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+
+MODEL_NAME=$1
+BACKEND=$2 # "mlm", "causal", "mntp", "enc_dec_mask", "enc_dec_prefix
+TRACK=$3 # "strict-small", "multimodal", "strict", "interaction"
+
+python -m evaluation_pipeline.collate_preds --model_path_or_name=$MODEL_NAME --backend=$BACKEND --fast --track=$TRACK

babylm_eval_repo/strict/scripts/download_evals.py

@@ -0,0 +1,7 @@
+from huggingface_hub import snapshot_download
+
+snapshot_download(
+    repo_id="BabyLM-community/BabyLM-2026-Strict-Evals",
+    repo_type="dataset",
+    local_dir=".",
+)

babylm_eval_repo/strict/scripts/eval_aoa.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+MODEL_NAME=$1
+BACKEND=$2
+TRACK_NAME=$3
+WORD_PATH=${4:-"evaluation_data/full_eval/aoa/cdi_childes.json"}
+OUTPUT_DIR=${5:-"results"}
+
+# please check evaluation_pipeline/AoA_word/README.md for more information on parameter setting
+
+# Set default parameters
+MIN_CONTEXT=${MIN_CONTEXT:-20}
+echo "Running AoA evaluation for model: $MODEL_NAME"
+echo "Word path: $WORD_PATH"
+echo "Output directory: $OUTPUT_DIR"
+# Create output directory if it doesn't exist
+mkdir -p $OUTPUT_DIR
+# Run the main evaluation
+python -m evaluation_pipeline.AoA_word.run \
+--model_name $MODEL_NAME \
+--backend $BACKEND \
+--track_name $TRACK_NAME \
+--word_path $WORD_PATH \
+--output_dir $OUTPUT_DIR 
+

babylm_eval_repo/strict/scripts/eval_devbench.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+
+MODEL_PATH=$1
+MODEL_TYPE=$2
+IMAGE_MODEL={$3:-$MODEL_PATH}
+
+# Supported MODEL_TYPE values:
+# git, flamingo, llava, flava, clip, blip, siglip, bridgetower, vilt, cvcl
+
+# If you need a different MODEL_TYPE, implement it in the `evaluation_pipeline/devbench/model_classes` folder.
+# (See other files in that folder for examples.)
+# Then add a wrapper to `evaluation_pipeline/devbench/eval.py`.
+# Be sure to submit a pull request so others can benefit from your implementation!
+
+python -m evaluation_pipeline.devbench.eval --model $MODEL_PATH \
+    --model_type $MODEL_TYPE \
+    --image_model $IMAGE_MODEL

babylm_eval_repo/strict/scripts/eval_finetuning.sh

@@ -0,0 +1,174 @@
+#!/bin/bash
+
+set -a
+source ../.env
+set +a
+
+# Default values
+MODEL_PATH=""
+LR=3e-5
+BSZ=32
+BIG_BSZ=16
+MAX_EPOCHS=10
+WSC_EPOCHS=30
+SEED=42
+WANDB_FLAG=""
+
+# Parse named arguments
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --model_path)
+            MODEL_PATH="$2"
+            shift 2
+            ;;
+        --lr)
+            LR="$2"
+            shift 2
+            ;;
+        --bsz)
+            BSZ="$2"
+            shift 2
+            ;;
+        --big_bsz)
+            BIG_BSZ="$2"
+            shift 2
+            ;;
+        --max_epochs)
+            MAX_EPOCHS="$2"
+            shift 2
+            ;;
+        --wsc_epochs)
+            WSC_EPOCHS="$2"
+            shift 2
+            ;;
+        --seed)
+            SEED="$2"
+            shift 2
+            ;;
+        --wandb)
+            WANDB_FLAG="--wandb"
+            shift
+            ;;
+        *)
+            echo "Unknown argument: $1"
+            exit 1
+            ;;
+    esac
+done
+
+for task in {boolq,multirc}; do
+    echo $task
+
+    python -m evaluation_pipeline.finetune.run \
+        --model_name_or_path "$MODEL_PATH" \
+        --train_data "evaluation_data/full_eval/glue_filtered/$task.train.jsonl" \
+        --valid_data "evaluation_data/full_eval/glue_filtered/$task.valid.jsonl" \
+        --predict_data "evaluation_data/full_eval/glue_filtered/$task.valid.jsonl" \
+        --task "$task" \
+        --num_labels 2 \
+        --batch_size $BIG_BSZ \
+        --learning_rate $LR \
+        --num_epochs $MAX_EPOCHS \
+        --sequence_length 512 \
+        --results_dir "results" \
+        --save \
+        --save_dir "models" \
+        --metrics accuracy f1 mcc \
+        --metric_for_valid accuracy \
+        --seed $SEED \
+        --verbose \
+        --padding_side left \
+        --take_final \
+        $WANDB_FLAG
+done
+
+python -m evaluation_pipeline.finetune.run \
+    --model_name_or_path "$MODEL_PATH" \
+    --train_data "evaluation_data/full_eval/glue_filtered/rte.train.jsonl" \
+    --valid_data "evaluation_data/full_eval/glue_filtered/rte.valid.jsonl" \
+    --predict_data "evaluation_data/full_eval/glue_filtered/rte.valid.jsonl" \
+    --task rte \
+    --num_labels 2 \
+    --batch_size $BSZ \
+    --learning_rate $LR \
+    --num_epochs $MAX_EPOCHS \
+    --sequence_length 512 \
+    --results_dir "results" \
+    --save \
+    --save_dir "models" \
+    --metrics accuracy f1 mcc \
+    --metric_for_valid accuracy \
+    --seed $SEED \
+    --verbose \
+    --padding_side left \
+    --take_final \
+    $WANDB_FLAG
+
+python -m evaluation_pipeline.finetune.run \
+    --model_name_or_path "$MODEL_PATH" \
+    --train_data "evaluation_data/full_eval/glue_filtered/wsc.train.jsonl" \
+    --valid_data "evaluation_data/full_eval/glue_filtered/wsc.valid.jsonl" \
+    --predict_data "evaluation_data/full_eval/glue_filtered/wsc.valid.jsonl" \
+    --task wsc \
+    --num_labels 2 \
+    --batch_size $BSZ \
+    --learning_rate $LR \
+    --num_epochs $WSC_EPOCHS \
+    --sequence_length 512 \
+    --results_dir "results" \
+    --save \
+    --save_dir "models" \
+    --metrics accuracy f1 mcc \
+    --metric_for_valid accuracy \
+    --seed $SEED \
+    --verbose \
+    --padding_side left \
+    --take_final \
+    $WANDB_FLAG
+
+for task in {mrpc,qqp}; do
+        
+    python -m evaluation_pipeline.finetune.run \
+        --model_name_or_path "$MODEL_PATH" \
+        --train_data "evaluation_data/full_eval/glue_filtered/$task.train.jsonl" \
+        --valid_data "evaluation_data/full_eval/glue_filtered/$task.valid.jsonl" \
+        --predict_data "evaluation_data/full_eval/glue_filtered/$task.valid.jsonl" \
+        --task "$task" \
+        --num_labels 2 \
+        --batch_size $BSZ \
+        --learning_rate $LR \
+        --num_epochs $MAX_EPOCHS \
+        --sequence_length 512 \
+        --results_dir "results" \
+        --save \
+        --save_dir "models" \
+        --metrics accuracy f1 mcc \
+        --metric_for_valid f1 \
+        --seed $SEED \
+        --verbose \
+	--padding_side left \
+	--take_final \
+        $WANDB_FLAG
+done
+
+python -m evaluation_pipeline.finetune.run \
+    --model_name_or_path "$MODEL_PATH" \
+    --train_data "evaluation_data/full_eval/glue_filtered/mnli.train.jsonl" \
+    --valid_data "evaluation_data/full_eval/glue_filtered/mnli.valid.jsonl" \
+    --predict_data "evaluation_data/full_eval/glue_filtered/mnli.valid.jsonl" \
+    --task mnli \
+    --num_labels 3 \
+    --batch_size $BSZ \
+    --learning_rate $LR \
+    --num_epochs $MAX_EPOCHS \
+    --sequence_length 512 \
+    --results_dir "results" \
+    --save \
+    --save_dir "models" \
+    --metrics accuracy \
+    --metric_for_valid accuracy \
+    --seed $SEED \
+    --verbose \
+    --padding_side left \
+    --take_final \
+    $WANDB_FLAG

babylm_eval_repo/strict/scripts/eval_multimodal.sh

@@ -0,0 +1,35 @@
+#!/bin/bash
+
+MODEL_PATH=$1
+BACKEND=$2
+MODEL_TYPE=$3
+IMAGE_MODEL=$4
+EVAL_DIR=${5:-"evaluation_data/full_eval"}
+
+if [[ "$BACKEND" == *"enc_dec"* ]]; then
+    BACKEND_READ="enc_dec"
+else
+    BACKEND_READ=$BACKEND
+fi
+
+echo $BACKEND_READ
+
+# Winoground and VQA
+python -m evaluation_pipeline.sentence_zero_shot.run --model_path_or_name $MODEL_PATH --backend $BACKEND --task vqa --data_path "evaluation_data/full_eval/vqa_filtered" --save_predictions --images_path=HuggingFaceM4/VQAv2 --image_split=validation --batch_size=1
+python -m evaluation_pipeline.sentence_zero_shot.run --model_path_or_name $MODEL_PATH --backend $BACKEND --task winoground --data_path "evaluation_data/full_eval/winoground_filtered" --save_predictions --images_path=facebook/winoground --image_split=test --batch_size=1
+
+# Devbench
+# Supported MODEL_TYPE values:
+# git, flamingo, llava, flava, clip, blip, siglip, bridgetower, vilt, cvcl
+
+# If you need a different MODEL_TYPE, implement it in the `evaluation_pipeline/devbench/model_classes` folder.
+# (See other files in that folder for examples.)
+# Then add a wrapper to `evaluation_pipeline/devbench/eval.py`.
+# Be sure to submit a pull request so others can benefit from your implementation!
+
+python -m evaluation_pipeline.devbench.eval --model $MODEL_PATH \
+    --model_type $MODEL_TYPE \
+    --image_model $IMAGE_MODEL
+
+
+

babylm_eval_repo/strict/scripts/eval_zero_shot.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+MODEL_PATH=$1
+BACKEND=$2
+EVAL_DIR=${3:-"evaluation_data/full_eval"}
+REVISION=${4:-main}
+
+if [[ "$BACKEND" == *"enc_dec"* ]]; then
+    BACKEND_READ="enc_dec"
+else
+    BACKEND_READ=$BACKEND
+fi
+
+echo $BACKEND_READ
+
+python -m evaluation_pipeline.sentence_zero_shot.run --model_path_or_name $MODEL_PATH --backend $BACKEND --task blimp --data_path "${EVAL_DIR}/blimp_filtered" --save_predictions --revision_name $REVISION
+python -m evaluation_pipeline.sentence_zero_shot.run --model_path_or_name $MODEL_PATH --backend $BACKEND --task blimp --data_path "${EVAL_DIR}/supplement_filtered" --save_predictions --revision_name $REVISION
+python -m evaluation_pipeline.sentence_zero_shot.run --model_path_or_name $MODEL_PATH --backend $BACKEND --task ewok --data_path "${EVAL_DIR}/ewok_filtered" --save_predictions --revision_name $REVISION
+python -m evaluation_pipeline.sentence_zero_shot.run --model_path_or_name $MODEL_PATH --backend $BACKEND --task entity_tracking --data_path "${EVAL_DIR}/entity_tracking" --save_predictions --revision_name $REVISION
+python -m evaluation_pipeline.sentence_zero_shot.run --model_path_or_name $MODEL_PATH --backend $BACKEND --task comps --data_path "${EVAL_DIR}/comps" --save_predictions --revision_name $REVISION
+python -m evaluation_pipeline.reading.run --model_path_or_name $MODEL_PATH --backend $BACKEND_READ --data_path "${EVAL_DIR}/reading/reading_data.csv" --revision_name $REVISION

babylm_eval_repo/strict/scripts/eval_zero_shot_fast.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+
+MODEL_PATH=$1
+REVISION_NAME=$2
+BACKEND=$3
+EVAL_DIR=${4:-"evaluation_data/fast_eval"}
+
+if [[ "$BACKEND" == *"enc_dec"* ]]; then
+    BACKEND_READ="enc_dec"
+else
+    BACKEND_READ=$BACKEND
+fi
+
+python -m evaluation_pipeline.sentence_zero_shot.run --model_path_or_name $MODEL_PATH --backend $BACKEND --task blimp --data_path "${EVAL_DIR}/blimp_fast" --save_predictions --revision_name $REVISION_NAME
+python -m evaluation_pipeline.sentence_zero_shot.run --model_path_or_name $MODEL_PATH --backend $BACKEND --task blimp --data_path "${EVAL_DIR}/supplement_fast" --save_predictions --revision_name $REVISION_NAME
+python -m evaluation_pipeline.sentence_zero_shot.run --model_path_or_name $MODEL_PATH --backend $BACKEND --task ewok --data_path "${EVAL_DIR}/ewok_fast" --save_predictions --revision_name $REVISION_NAME
+python -m evaluation_pipeline.sentence_zero_shot.run --model_path_or_name $MODEL_PATH --backend $BACKEND --task entity_tracking --data_path "${EVAL_DIR}/entity_tracking_fast" --save_predictions --revision_name $REVISION_NAME
+python -m evaluation_pipeline.reading.run --model_path_or_name $MODEL_PATH --backend $BACKEND_READ --data_path "${EVAL_DIR}/reading/reading_data.csv" --revision_name $REVISION_NAME

babylm_eval_repo/strict/scripts/eval_zero_shot_fast_all_revisions.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+
+MODEL_PATH=$1
+BACKEND=$2
+TRACK=$3
+EVAL_DIR=${4:-"evaluation_data/fast_eval"}
+
+for i in {1..9}; do
+    checkpoint="chck_${i}M"
+    echo "Evaluating checkpoint ${checkpoint}"
+    bash scripts/eval_zero_shot_fast.sh $MODEL_PATH $checkpoint $BACKEND $EVAL_DIR
+done
+
+for i in {10..100..10}; do
+    checkpoint="chck_${i}M"
+    echo "Evaluating checkpoint ${checkpoint}"
+    bash scripts/eval_zero_shot_fast.sh $MODEL_PATH $checkpoint $BACKEND $EVAL_DIR
+done
+
+# Conditional on whether the track is strict-small
+if [[ "$3" != "strict-small" ]]; then
+    for i in {200..1000..100}; do
+	checkpoint="chck_${i}M"
+	echo "Evaluating checkpoint ${checkpoint}"
+	bash scripts/eval_zero_shot_fast.sh $MODEL_PATH $checkpoint $BACKEND $EVAL_DIR
+    done
+fi

babylm_eval_repo/strict/scripts/eval_zero_shot_global_piqa.sh

@@ -0,0 +1,68 @@
+#!/bin/bash
+# GlobalPIQA-only zero-shot evaluation for a single model, covering both the
+# main checkpoint and all intermediate (fast-revision) checkpoints. This is a
+# combined, GlobalPIQA-dedicated version of eval_zero_shot.sh +
+# eval_zero_shot_fast_all_revisions.sh.
+#
+# GlobalPIQA is not subsampled for the fast evaluation, so the parallel/
+# nonparallel data are identical under full_eval and fast_eval. The main
+# checkpoint (which feeds the full-eval collation) is scored against full_eval;
+# the intermediate checkpoints (which feed the fast-eval collation) are scored
+# against fast_eval, mirroring every other fast task. Predictions are
+# length-normalized at inference time, so the resulting accuracy is acc_norm.
+#
+# Usage (run from the strict/ directory):
+#   bash scripts/eval_zero_shot_global_piqa.sh MODEL_PATH BACKEND TRACK [EVAL_DIR] [FAST_EVAL_DIR]
+#
+#   MODEL_PATH     HuggingFace repo id or local path to the model.
+#   BACKEND        Evaluation backend (causal, mlm, mntp, enc_dec_mask, enc_dec_prefix).
+#   TRACK          strict | strict-small (strict-small stops at the 100M checkpoint).
+#   EVAL_DIR       Optional full-eval data dir for the main checkpoint (default evaluation_data/full_eval).
+#   FAST_EVAL_DIR  Optional fast-eval data dir for intermediate checkpoints (default evaluation_data/fast_eval).
+
+MODEL_PATH=$1
+BACKEND=$2
+TRACK=$3
+EVAL_DIR=${4:-"evaluation_data/full_eval"}
+FAST_EVAL_DIR=${5:-"evaluation_data/fast_eval"}
+
+# Run the two GlobalPIQA tasks for one checkpoint against the given data dir.
+# With no revision argument the pipeline defaults to the "main" checkpoint;
+# otherwise the given revision name is passed through.
+run_global_piqa () {
+    local revision=$1
+    local data_dir=$2
+    local revision_args=""
+    if [[ -n "$revision" ]]; then
+        revision_args="--revision_name $revision"
+    fi
+
+    python -m evaluation_pipeline.sentence_zero_shot.run --model_path_or_name $MODEL_PATH --backend $BACKEND --task global_piqa_parallel --data_path "${data_dir}/global_piqa_parallel" --save_predictions $revision_args
+    python -m evaluation_pipeline.sentence_zero_shot.run --model_path_or_name $MODEL_PATH --backend $BACKEND --task global_piqa_nonparallel --data_path "${data_dir}/global_piqa_nonparallel" --save_predictions $revision_args
+}
+
+# Main checkpoint -> full_eval
+echo "Evaluating checkpoint main"
+run_global_piqa "" "$EVAL_DIR"
+
+# Intermediate checkpoints -> fast_eval: 1M..9M, then 10M..100M
+for i in {1..9}; do
+    checkpoint="chck_${i}M"
+    echo "Evaluating checkpoint ${checkpoint}"
+    run_global_piqa $checkpoint "$FAST_EVAL_DIR"
+done
+
+for i in {10..100..10}; do
+    checkpoint="chck_${i}M"
+    echo "Evaluating checkpoint ${checkpoint}"
+    run_global_piqa $checkpoint "$FAST_EVAL_DIR"
+done
+
+# Full strict track continues past 100M up to 1000M; strict-small stops at 100M.
+if [[ "$TRACK" != "strict-small" ]]; then
+    for i in {200..1000..100}; do
+        checkpoint="chck_${i}M"
+        echo "Evaluating checkpoint ${checkpoint}"
+        run_global_piqa $checkpoint "$FAST_EVAL_DIR"
+    done
+fi

babylm_eval_repo/strict/scripts/print_results_table.py

@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+"""Print aggregate markdown tables of strict-track results from results/.
+
+Zero-shot results are read from:
+  results/{model}/main/{eval_type}/causal/{task}/{subtask}/best_temperature_report.txt
+
+Finetune results are read from:
+  results/{model}/main/finetune/{task}/results.txt
+"""
+
+from pathlib import Path
+
+# Primary metric per GLUE task, matching the leaderboard definition in about.py.
+FINETUNE_METRIC = {
+    "boolq":   "accuracy",
+    "mnli":    "accuracy",
+    "mrpc":    "f1",
+    "multirc": "accuracy",
+    "qqp":     "f1",
+    "rte":     "accuracy",
+    "wsc":     "accuracy",
+}
+
+
+def parse_average_score(report_path: Path) -> float | None:
+    """Return the average score from a best_temperature_report.txt file.
+
+    Handles both '### AVERAGE ACCURACY' and '### AVERAGE SPEARMAN'S RHO' headers.
+    """
+    lines = report_path.read_text().splitlines()
+    for i, line in enumerate(lines):
+        if line.strip().startswith("### AVERAGE"):
+            if i + 1 < len(lines):
+                try:
+                    return float(lines[i + 1].strip())
+                except ValueError:
+                    return None
+    return None
+
+
+def parse_finetune_score(results_path: Path, metric: str) -> float | None:
+    """Return the value of `metric` from a finetune results.txt file."""
+    for line in results_path.read_text().splitlines():
+        key, _, value = line.partition(":")
+        if key.strip() == metric:
+            try:
+                return float(value.strip()) * 100
+            except ValueError:
+                return None
+    return None
+
+
+def fmt_row(cells: list[str]) -> str:
+    return "| " + " | ".join(cells) + " |"
+
+
+def print_table(
+    all_models: dict[str, dict[str, float]],
+    row_order: list[str],
+    models: list[str],
+) -> None:
+    print(fmt_row(["task"] + models))
+    print(fmt_row(["---"] + ["---"] * len(models)))
+    for row_key in row_order:
+        vals = [all_models[model].get(row_key) for model in models]
+        best = max((v for v in vals if v is not None), default=None)
+        row = [row_key]
+        for val in vals:
+            if val is None:
+                row.append("")
+            elif val == best:
+                row.append(f"**{val:.2f}**")
+            else:
+                row.append(f"{val:.2f}")
+        print(fmt_row(row))
+
+
+def main():
+    results_dir = Path(__file__).parent.parent / "results"
+
+    # --- Zero-shot ---
+    zeroshot_models: dict[str, dict[str, float]] = {}
+    zeroshot_row_order: list[str] = []
+
+    for report in sorted(results_dir.glob("*/main/*/causal/*/*/best_temperature_report.txt")):
+        # parts: results / model / main / eval_type / causal / task / subtask / filename
+        parts = report.parts
+        model     = parts[-7]
+        eval_type = parts[-5]
+        task      = parts[-3]
+        subtask   = parts[-2]
+
+        score = parse_average_score(report)
+        if score is None:
+            continue
+
+        row_key = f"{eval_type}/{task}/{subtask}"
+        if row_key not in zeroshot_row_order:
+            zeroshot_row_order.append(row_key)
+        zeroshot_models.setdefault(model, {})[row_key] = score
+
+    # --- Finetune ---
+    finetune_models: dict[str, dict[str, float]] = {}
+    finetune_row_order: list[str] = []
+
+    for results_file in sorted(results_dir.glob("*/main/finetune/*/results.txt")):
+        # parts: results / model / main / finetune / task / filename
+        parts = results_file.parts
+        model = parts[-5]
+        task  = parts[-2]
+
+        metric = FINETUNE_METRIC.get(task, "accuracy")
+        score = parse_finetune_score(results_file, metric)
+        if score is None:
+            continue
+
+        row_key = f"{task} ({metric})"
+        if row_key not in finetune_row_order:
+            finetune_row_order.append(row_key)
+        finetune_models.setdefault(model, {})[row_key] = score
+
+    if not zeroshot_models and not finetune_models:
+        print("No results found.")
+        return
+
+    all_models = sorted(set(list(zeroshot_models) + list(finetune_models)))
+
+    if zeroshot_models:
+        print("## Zero-shot\n")
+        print_table(zeroshot_models, zeroshot_row_order, all_models)
+
+    if finetune_models:
+        print("\n## Finetune (GLUE)\n")
+        print_table(finetune_models, finetune_row_order, all_models)
+
+
+if __name__ == "__main__":
+    main()

checkpoints/xpertgpt_fresh/chck_100M/config.json

@@ -0,0 +1,17 @@
+{
+  "auto_map": {
+    "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
+    "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
+    "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
+  },
+  "vocab_size": 16384,
+  "block_size": 512,
+  "d_model": 256,
+  "hidden_size": 256,
+  "d_thin": 384,
+  "num_layers": 6,
+  "num_blocks": 4,
+  "capacity_factor": 0.5,
+  "dropout": 0.1,
+  "model_type": "xpertgpt"
+}

checkpoints/xpertgpt_fresh/chck_100M/configuration_xpertgpt.py

@@ -0,0 +1,36 @@
+from transformers import PretrainedConfig
+
+class XpertGPTConfig(PretrainedConfig):
+    model_type = "xpertgpt"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 256,
+        d_thin: int = 384,
+        num_layers: int = 6,
+        num_blocks: int = 4,
+        capacity_factor: float = 0.5,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+        self.num_hidden_layers = num_layers
+
+        super().__init__(**kwargs)

checkpoints/xpertgpt_fresh/chck_100M/modeling_xpertgpt.py

@@ -0,0 +1,456 @@
+import os
+import math
+import random
+import inspect
+from typing import Optional, Tuple, Dict, Any
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin, AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+# ─────────────────────────────────────────────────────────────
+# Configuration Classes
+# ─────────────────────────────────────────────────────────────
+
+@dataclass
+class XpertGPTModelConfig:
+    vocab_size: int = 16384
+    block_size: int = 512
+    d_model: int = 256
+    d_thin: int = 384
+    num_layers: int = 6
+    num_blocks: int = 4
+    capacity_factor: float = 0.5
+    dropout: float = 0.1
+
+class XpertGPTConfig(PretrainedConfig):
+    model_type = "xpertgpt"
+    auto_map = {
+        "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
+        "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
+        "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 256,
+        d_thin: int = 384,
+        num_layers: int = 6,
+        num_blocks: int = 4,
+        capacity_factor: float = 0.5,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+        self.num_hidden_layers = num_layers
+
+        super().__init__(**kwargs)
+
+# ─────────────────────────────────────────────────────────────
+# ROPE HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def _precompute_rope_freqs(head_dim: int, seq_len: int, device: torch.device, theta: float = 10000.0):
+    assert head_dim % 2 == 0, "head_dim must be divisible by 2 for RoPE"
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos(), emb.sin()
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    L = x.size(2)
+    cos = cos[:L, :].unsqueeze(0).unsqueeze(1) 
+    sin = sin[:L, :].unsqueeze(0).unsqueeze(1) 
+    
+    half_dim = x.size(-1) // 2
+    x1 = x[..., :half_dim]
+    x2 = x[..., half_dim:]
+    rotated_x = torch.cat((-x2, x1), dim=-1)
+    
+    return (x * cos) + (rotated_x * sin)
+
+# ─────────────────────────────────────────────────────────────
+# 1.  SLIDING WINDOW ATTENTION
+# ─────────────────────────────────────────────────────────────
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size=None):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.num_heads  = num_heads
+        self.window_size = window_size
+        self.head_dim   = dim // num_heads
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        B, L, D = x.size()
+
+        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            past_len = past_k.size(2)
+            q_cos, q_sin = _precompute_rope_freqs(self.head_dim, past_len + L, x.device)
+            q = _apply_rope(q, q_cos[past_len:, :], q_sin[past_len:, :])
+            k = _apply_rope(k, q_cos[past_len:, :], q_sin[past_len:, :])
+            
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        else:
+            cos, sin = _precompute_rope_freqs(self.head_dim, L, x.device)
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+
+        L_kv = k.size(2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        if bidirectional:
+            if self.window_size is not None:
+                past_len = L_kv - L
+                pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+                pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+                dist  = torch.abs(pos_i - pos_j)
+                win_mask = dist < self.window_size
+                scores = scores.masked_fill(
+                    ~win_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+                )
+        else:
+            past_len = L_kv - L
+            pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+            pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+            dist  = pos_i - pos_j
+
+            causal_mask = dist >= 0
+            if self.window_size is not None:
+                causal_mask = causal_mask & (dist < self.window_size)
+
+            scores = scores.masked_fill(
+                ~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+            )
+
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn, v)                                
+        out = out.transpose(1, 2).contiguous().view(B, L, D)
+        out = self.o_proj(out)
+
+        if use_cache:
+            if self.window_size is not None:
+                present_kv = (
+                    k[:, :, -self.window_size:, :],
+                    v[:, :, -self.window_size:, :]
+                )
+            else:
+                present_kv = (k, v)
+        else:
+            present_kv = None
+
+        return out, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 2.  MSIT BRANCH BLOCK (Pre-Norm)
+# ─────────────────────────────────────────────────────────────
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        # Keep parameter count comparable to regular FFN with dim * 4:
+        # 3 * dim * hidden_dim ~= 8 * dim^2  =>  hidden_dim = 8/3 * dim
+        hidden_dim = int(dim * 8 / 3)
+        hidden_dim = ((hidden_dim + 7) // 8) * 8
+        self.fc1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.fc2 = nn.Linear(dim, hidden_dim, bias=False)
+        self.fc3 = nn.Linear(hidden_dim, dim, bias=False)
+
+    def forward(self, x):
+        return self.fc3(F.silu(self.fc1(x)) * self.fc2(x))
+
+class MSITBranchBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size):
+        super().__init__()
+        self.ln1  = nn.LayerNorm(dim)
+        self.attn = SlidingWindowAttention(dim, num_heads, window_size)
+        self.ln2  = nn.LayerNorm(dim)
+        self.ffn  = SwiGLU(dim)
+        self.ln3  = nn.LayerNorm(dim)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        attn_out, present_kv = self.attn(
+            self.ln1(x), past_kv, use_cache, bidirectional
+        )
+        x = x + attn_out
+        x = x + self.ffn(self.ln2(x))
+        x = self.ln3(x)
+        return x, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 3.  MoEP-MSIT ARCHITECTURE BLOCK  (Expert Choice routing)
+# ─────────────────────────────────────────────────────────────
+
+class MoEPMSITBlock(nn.Module):
+    def __init__(self, d_model: int = 512, d_thin: int = 192, num_blocks: int = 14,
+                 capacity_factor: float = 0.5):
+        super().__init__()
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+
+        # 1. Global Block (Dense, d_model)
+        num_heads_global = max(1, d_model // 64)
+        self.global_block = MSITBranchBlock(d_model, num_heads_global, window_size=None)
+
+        # 2. Router
+        self.router_ln = nn.LayerNorm(d_model)
+        self.w_router = nn.Linear(d_model, num_blocks, bias=False)
+
+        # 3. Shrink Projection
+        self.w_down = nn.Linear(d_model, d_thin, bias=False)
+
+        # 4. Thin Parallel Blocks (d_thin)
+        self.windows = [64, 16, 8, 4] + [None] * (num_blocks - 4)
+        self.heads = [max(1, d_thin // 64)] * num_blocks  
+        
+        self.thin_blocks = nn.ModuleList([
+            MSITBranchBlock(d_thin, self.heads[i], self.windows[i])
+            for i in range(num_blocks)
+        ])
+
+        # 6. Grow Projection
+        self.w_up = nn.Linear(d_thin, d_model, bias=False)
+        self.last_topk_indices = None
+        self.ln_post_moe = nn.LayerNorm(d_model)
+
+    def forward(self, x_0: torch.Tensor, past_kvs=None, use_cache: bool = False, bidirectional: bool = False):
+        B, T, D = x_0.size()
+        n_tokens = B * T
+
+        # Step 1: Global Block
+        pkv_g = past_kvs[0] if past_kvs else None
+        x_1, nkv_g = self.global_block(x_0, pkv_g, use_cache, bidirectional)
+
+        # Step 2: Gated input stream
+        x_2 = x_1
+
+        # Router scores
+        r_logits = self.w_router(self.router_ln(x_2)).view(n_tokens, self.num_blocks)
+        r_probs = F.softmax(r_logits, dim=-1)
+
+        # Per-expert capacity: k = (n * c) / e
+        k_capacity = max(1, int(round(n_tokens * self.capacity_factor / self.num_blocks)))
+        k_capacity = min(k_capacity, n_tokens)
+
+        # Expert Choice routing: topk over the token axis for each expert
+        expert_token_scores = r_probs.transpose(0, 1) # (num_blocks, n_tokens)
+        topk_scores, topk_token_idx = torch.topk(expert_token_scores, k_capacity, dim=-1)
+        self.last_topk_indices = topk_token_idx
+
+        # Load balancing is guaranteed by construction in Expert Choice
+        layer_aux_loss = x_2.new_zeros(())
+
+        # Shrink projection
+        x_2_thin = self.w_down(x_2) # (B, T, d_thin)
+        x_2_thin_flat = x_2_thin.view(n_tokens, self.d_thin)
+
+        # Expert computations
+        new_kvs = [nkv_g]
+        expert_outputs_flat = torch.zeros(n_tokens, self.d_model, device=x_2.device, dtype=x_2.dtype)
+
+        for i, block in enumerate(self.thin_blocks):
+            sel_idx = topk_token_idx[i]
+            bucket_in = x_2_thin_flat[sel_idx].unsqueeze(0) # (1, k_capacity, d_thin)
+
+            pkv_i = past_kvs[i + 1] if past_kvs else None
+            bucket_out, nkv_i = block(bucket_in, pkv_i, use_cache, bidirectional)
+            if use_cache:
+                new_kvs.append(nkv_i)
+
+            bucket_out = bucket_out.squeeze(0) # (k_capacity, d_thin)
+            bucket_out_full = self.w_up(bucket_out) # (k_capacity, d_model)
+
+            gate = topk_scores[i].unsqueeze(-1) # (k_capacity, 1)
+            expert_outputs_flat.index_add_(0, sel_idx, bucket_out_full * gate)
+
+        x_3_full = expert_outputs_flat.view(B, T, self.d_model)
+        out = x_2 + x_3_full
+        out = self.ln_post_moe(out)
+
+        present_kvs = tuple(new_kvs) if use_cache else None
+        return out, present_kvs, layer_aux_loss
+
+# ─────────────────────────────────────────────────────────────
+# 4.  RAW XpertGPT MODEL
+# ─────────────────────────────────────────────────────────────
+
+class XpertGPTModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.wte      = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop_emb = nn.Dropout(cfg.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(
+                d_model=cfg.d_model, 
+                d_thin=cfg.d_thin, 
+                num_blocks=cfg.num_blocks, 
+                capacity_factor=cfg.capacity_factor
+            )
+            for _ in range(cfg.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(cfg.d_model)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.wte.weight = self.lm_head.weight
+
+    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None, bidirectional: bool = False):
+        x = self.drop_emb(self.wte(input_ids))
+        total_aux_loss = 0.0
+        
+        for block in self.blocks:
+            x, _, layer_aux = block(x, past_kvs=None, use_cache=False, bidirectional=bidirectional)
+            total_aux_loss += layer_aux
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            ce_loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), targets.view(-1), ignore_index=-100)
+            avg_aux_loss = total_aux_loss / self.cfg.num_layers
+            loss = ce_loss + (0.01 * avg_aux_loss)
+
+        return logits, loss
+
+# ─────────────────────────────────────────────────────────────
+# 5.  HUGGING FACE MODEL WRAPPERS
+# ─────────────────────────────────────────────────────────────
+
+class XpertGPTModelWrapper(PreTrainedModel):
+    config_class = XpertGPTConfig
+    base_model_prefix = "transformer"
+
+    def __init__(self, config: XpertGPTConfig):
+        super().__init__(config)
+        self.wte      = nn.Embedding(config.vocab_size, config.d_model)
+        self.drop_emb = nn.Dropout(config.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(config.d_model, config.d_thin, config.num_blocks, config.capacity_factor)
+            for _ in range(config.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(config.d_model)
+        self.post_init()
+
+    def forward(self, input_ids, **kwargs):
+        x = self.drop_emb(self.wte(input_ids))
+        for block in self.blocks:
+            x, _, _ = block(x, past_kvs=None, use_cache=False, bidirectional=False)
+        x = self.ln_f(x)
+        return BaseModelOutputWithPast(last_hidden_state=x)
+
+
+class XpertGPTForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = XpertGPTConfig
+    base_model_prefix = "transformer"
+    _no_split_modules = ["MoEPMSITBlock"]
+    _tied_weights_keys = {"transformer.lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config: XpertGPTConfig):
+        super().__init__(config)
+        self.transformer = XpertGPTModelWrapper(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.post_init()
+        
+        # State-dict pre-hook for backwards compatibility with checkpoint key naming
+        def _prefix_cleaner(state_dict, prefix, local_metadata, Moore, missing_keys, unexpected_keys, error_msgs):
+            keys = list(state_dict.keys())
+            for k in keys:
+                if k.startswith("transformer."):
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                elif f"{prefix}transformer." in k:
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                    
+        self._register_load_state_dict_pre_hook(_prefix_cleaner)
+
+    def tie_weights(self, **kwargs):
+        if hasattr(self, "transformer") and hasattr(self.transformer, "wte") and hasattr(self.transformer, "lm_head"):
+            self.transformer.wte.weight = self.lm_head.weight
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(self,
+                input_ids: Optional[torch.LongTensor] = None,
+                attention_mask: Optional[torch.FloatTensor] = None,
+                labels: Optional[torch.LongTensor] = None,
+                **kwargs) -> CausalLMOutputWithPast:
+        outputs = self.transformer(input_ids)
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# Register with auto-mapping
+AutoConfig.register("xpertgpt", XpertGPTConfig)
+AutoModel.register(XpertGPTConfig, XpertGPTModelWrapper)
+AutoModelForCausalLM.register(XpertGPTConfig, XpertGPTForCausalLM)

checkpoints/xpertgpt_fresh/chck_100M/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8d879766106dff337673cee6ceb268e706b04faefa2196e4ce93b1f1686e24f
+size 210797423

checkpoints/xpertgpt_fresh/chck_100M/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "[UNK]"
+}

checkpoints/xpertgpt_fresh/chck_10M/config.json

@@ -0,0 +1,17 @@
+{
+  "auto_map": {
+    "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
+    "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
+    "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
+  },
+  "vocab_size": 16384,
+  "block_size": 512,
+  "d_model": 256,
+  "hidden_size": 256,
+  "d_thin": 384,
+  "num_layers": 6,
+  "num_blocks": 4,
+  "capacity_factor": 0.5,
+  "dropout": 0.1,
+  "model_type": "xpertgpt"
+}

checkpoints/xpertgpt_fresh/chck_10M/configuration_xpertgpt.py

@@ -0,0 +1,36 @@
+from transformers import PretrainedConfig
+
+class XpertGPTConfig(PretrainedConfig):
+    model_type = "xpertgpt"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 256,
+        d_thin: int = 384,
+        num_layers: int = 6,
+        num_blocks: int = 4,
+        capacity_factor: float = 0.5,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+        self.num_hidden_layers = num_layers
+
+        super().__init__(**kwargs)

checkpoints/xpertgpt_fresh/chck_10M/modeling_xpertgpt.py

@@ -0,0 +1,456 @@
+import os
+import math
+import random
+import inspect
+from typing import Optional, Tuple, Dict, Any
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin, AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+# ─────────────────────────────────────────────────────────────
+# Configuration Classes
+# ─────────────────────────────────────────────────────────────
+
+@dataclass
+class XpertGPTModelConfig:
+    vocab_size: int = 16384
+    block_size: int = 512
+    d_model: int = 256
+    d_thin: int = 384
+    num_layers: int = 6
+    num_blocks: int = 4
+    capacity_factor: float = 0.5
+    dropout: float = 0.1
+
+class XpertGPTConfig(PretrainedConfig):
+    model_type = "xpertgpt"
+    auto_map = {
+        "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
+        "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
+        "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 256,
+        d_thin: int = 384,
+        num_layers: int = 6,
+        num_blocks: int = 4,
+        capacity_factor: float = 0.5,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+        self.num_hidden_layers = num_layers
+
+        super().__init__(**kwargs)
+
+# ─────────────────────────────────────────────────────────────
+# ROPE HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def _precompute_rope_freqs(head_dim: int, seq_len: int, device: torch.device, theta: float = 10000.0):
+    assert head_dim % 2 == 0, "head_dim must be divisible by 2 for RoPE"
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos(), emb.sin()
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    L = x.size(2)
+    cos = cos[:L, :].unsqueeze(0).unsqueeze(1) 
+    sin = sin[:L, :].unsqueeze(0).unsqueeze(1) 
+    
+    half_dim = x.size(-1) // 2
+    x1 = x[..., :half_dim]
+    x2 = x[..., half_dim:]
+    rotated_x = torch.cat((-x2, x1), dim=-1)
+    
+    return (x * cos) + (rotated_x * sin)
+
+# ─────────────────────────────────────────────────────────────
+# 1.  SLIDING WINDOW ATTENTION
+# ─────────────────────────────────────────────────────────────
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size=None):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.num_heads  = num_heads
+        self.window_size = window_size
+        self.head_dim   = dim // num_heads
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        B, L, D = x.size()
+
+        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            past_len = past_k.size(2)
+            q_cos, q_sin = _precompute_rope_freqs(self.head_dim, past_len + L, x.device)
+            q = _apply_rope(q, q_cos[past_len:, :], q_sin[past_len:, :])
+            k = _apply_rope(k, q_cos[past_len:, :], q_sin[past_len:, :])
+            
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        else:
+            cos, sin = _precompute_rope_freqs(self.head_dim, L, x.device)
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+
+        L_kv = k.size(2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        if bidirectional:
+            if self.window_size is not None:
+                past_len = L_kv - L
+                pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+                pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+                dist  = torch.abs(pos_i - pos_j)
+                win_mask = dist < self.window_size
+                scores = scores.masked_fill(
+                    ~win_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+                )
+        else:
+            past_len = L_kv - L
+            pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+            pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+            dist  = pos_i - pos_j
+
+            causal_mask = dist >= 0
+            if self.window_size is not None:
+                causal_mask = causal_mask & (dist < self.window_size)
+
+            scores = scores.masked_fill(
+                ~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+            )
+
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn, v)                                
+        out = out.transpose(1, 2).contiguous().view(B, L, D)
+        out = self.o_proj(out)
+
+        if use_cache:
+            if self.window_size is not None:
+                present_kv = (
+                    k[:, :, -self.window_size:, :],
+                    v[:, :, -self.window_size:, :]
+                )
+            else:
+                present_kv = (k, v)
+        else:
+            present_kv = None
+
+        return out, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 2.  MSIT BRANCH BLOCK (Pre-Norm)
+# ─────────────────────────────────────────────────────────────
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        # Keep parameter count comparable to regular FFN with dim * 4:
+        # 3 * dim * hidden_dim ~= 8 * dim^2  =>  hidden_dim = 8/3 * dim
+        hidden_dim = int(dim * 8 / 3)
+        hidden_dim = ((hidden_dim + 7) // 8) * 8
+        self.fc1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.fc2 = nn.Linear(dim, hidden_dim, bias=False)
+        self.fc3 = nn.Linear(hidden_dim, dim, bias=False)
+
+    def forward(self, x):
+        return self.fc3(F.silu(self.fc1(x)) * self.fc2(x))
+
+class MSITBranchBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size):
+        super().__init__()
+        self.ln1  = nn.LayerNorm(dim)
+        self.attn = SlidingWindowAttention(dim, num_heads, window_size)
+        self.ln2  = nn.LayerNorm(dim)
+        self.ffn  = SwiGLU(dim)
+        self.ln3  = nn.LayerNorm(dim)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        attn_out, present_kv = self.attn(
+            self.ln1(x), past_kv, use_cache, bidirectional
+        )
+        x = x + attn_out
+        x = x + self.ffn(self.ln2(x))
+        x = self.ln3(x)
+        return x, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 3.  MoEP-MSIT ARCHITECTURE BLOCK  (Expert Choice routing)
+# ─────────────────────────────────────────────────────────────
+
+class MoEPMSITBlock(nn.Module):
+    def __init__(self, d_model: int = 512, d_thin: int = 192, num_blocks: int = 14,
+                 capacity_factor: float = 0.5):
+        super().__init__()
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+
+        # 1. Global Block (Dense, d_model)
+        num_heads_global = max(1, d_model // 64)
+        self.global_block = MSITBranchBlock(d_model, num_heads_global, window_size=None)
+
+        # 2. Router
+        self.router_ln = nn.LayerNorm(d_model)
+        self.w_router = nn.Linear(d_model, num_blocks, bias=False)
+
+        # 3. Shrink Projection
+        self.w_down = nn.Linear(d_model, d_thin, bias=False)
+
+        # 4. Thin Parallel Blocks (d_thin)
+        self.windows = [64, 16, 8, 4] + [None] * (num_blocks - 4)
+        self.heads = [max(1, d_thin // 64)] * num_blocks  
+        
+        self.thin_blocks = nn.ModuleList([
+            MSITBranchBlock(d_thin, self.heads[i], self.windows[i])
+            for i in range(num_blocks)
+        ])
+
+        # 6. Grow Projection
+        self.w_up = nn.Linear(d_thin, d_model, bias=False)
+        self.last_topk_indices = None
+        self.ln_post_moe = nn.LayerNorm(d_model)
+
+    def forward(self, x_0: torch.Tensor, past_kvs=None, use_cache: bool = False, bidirectional: bool = False):
+        B, T, D = x_0.size()
+        n_tokens = B * T
+
+        # Step 1: Global Block
+        pkv_g = past_kvs[0] if past_kvs else None
+        x_1, nkv_g = self.global_block(x_0, pkv_g, use_cache, bidirectional)
+
+        # Step 2: Gated input stream
+        x_2 = x_1
+
+        # Router scores
+        r_logits = self.w_router(self.router_ln(x_2)).view(n_tokens, self.num_blocks)
+        r_probs = F.softmax(r_logits, dim=-1)
+
+        # Per-expert capacity: k = (n * c) / e
+        k_capacity = max(1, int(round(n_tokens * self.capacity_factor / self.num_blocks)))
+        k_capacity = min(k_capacity, n_tokens)
+
+        # Expert Choice routing: topk over the token axis for each expert
+        expert_token_scores = r_probs.transpose(0, 1) # (num_blocks, n_tokens)
+        topk_scores, topk_token_idx = torch.topk(expert_token_scores, k_capacity, dim=-1)
+        self.last_topk_indices = topk_token_idx
+
+        # Load balancing is guaranteed by construction in Expert Choice
+        layer_aux_loss = x_2.new_zeros(())
+
+        # Shrink projection
+        x_2_thin = self.w_down(x_2) # (B, T, d_thin)
+        x_2_thin_flat = x_2_thin.view(n_tokens, self.d_thin)
+
+        # Expert computations
+        new_kvs = [nkv_g]
+        expert_outputs_flat = torch.zeros(n_tokens, self.d_model, device=x_2.device, dtype=x_2.dtype)
+
+        for i, block in enumerate(self.thin_blocks):
+            sel_idx = topk_token_idx[i]
+            bucket_in = x_2_thin_flat[sel_idx].unsqueeze(0) # (1, k_capacity, d_thin)
+
+            pkv_i = past_kvs[i + 1] if past_kvs else None
+            bucket_out, nkv_i = block(bucket_in, pkv_i, use_cache, bidirectional)
+            if use_cache:
+                new_kvs.append(nkv_i)
+
+            bucket_out = bucket_out.squeeze(0) # (k_capacity, d_thin)
+            bucket_out_full = self.w_up(bucket_out) # (k_capacity, d_model)
+
+            gate = topk_scores[i].unsqueeze(-1) # (k_capacity, 1)
+            expert_outputs_flat.index_add_(0, sel_idx, bucket_out_full * gate)
+
+        x_3_full = expert_outputs_flat.view(B, T, self.d_model)
+        out = x_2 + x_3_full
+        out = self.ln_post_moe(out)
+
+        present_kvs = tuple(new_kvs) if use_cache else None
+        return out, present_kvs, layer_aux_loss
+
+# ─────────────────────────────────────────────────────────────
+# 4.  RAW XpertGPT MODEL
+# ─────────────────────────────────────────────────────────────
+
+class XpertGPTModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.wte      = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop_emb = nn.Dropout(cfg.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(
+                d_model=cfg.d_model, 
+                d_thin=cfg.d_thin, 
+                num_blocks=cfg.num_blocks, 
+                capacity_factor=cfg.capacity_factor
+            )
+            for _ in range(cfg.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(cfg.d_model)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.wte.weight = self.lm_head.weight
+
+    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None, bidirectional: bool = False):
+        x = self.drop_emb(self.wte(input_ids))
+        total_aux_loss = 0.0
+        
+        for block in self.blocks:
+            x, _, layer_aux = block(x, past_kvs=None, use_cache=False, bidirectional=bidirectional)
+            total_aux_loss += layer_aux
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            ce_loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), targets.view(-1), ignore_index=-100)
+            avg_aux_loss = total_aux_loss / self.cfg.num_layers
+            loss = ce_loss + (0.01 * avg_aux_loss)
+
+        return logits, loss
+
+# ─────────────────────────────────────────────────────────────
+# 5.  HUGGING FACE MODEL WRAPPERS
+# ─────────────────────────────────────────────────────────────
+
+class XpertGPTModelWrapper(PreTrainedModel):
+    config_class = XpertGPTConfig
+    base_model_prefix = "transformer"
+
+    def __init__(self, config: XpertGPTConfig):
+        super().__init__(config)
+        self.wte      = nn.Embedding(config.vocab_size, config.d_model)
+        self.drop_emb = nn.Dropout(config.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(config.d_model, config.d_thin, config.num_blocks, config.capacity_factor)
+            for _ in range(config.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(config.d_model)
+        self.post_init()
+
+    def forward(self, input_ids, **kwargs):
+        x = self.drop_emb(self.wte(input_ids))
+        for block in self.blocks:
+            x, _, _ = block(x, past_kvs=None, use_cache=False, bidirectional=False)
+        x = self.ln_f(x)
+        return BaseModelOutputWithPast(last_hidden_state=x)
+
+
+class XpertGPTForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = XpertGPTConfig
+    base_model_prefix = "transformer"
+    _no_split_modules = ["MoEPMSITBlock"]
+    _tied_weights_keys = {"transformer.lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config: XpertGPTConfig):
+        super().__init__(config)
+        self.transformer = XpertGPTModelWrapper(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.post_init()
+        
+        # State-dict pre-hook for backwards compatibility with checkpoint key naming
+        def _prefix_cleaner(state_dict, prefix, local_metadata, Moore, missing_keys, unexpected_keys, error_msgs):
+            keys = list(state_dict.keys())
+            for k in keys:
+                if k.startswith("transformer."):
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                elif f"{prefix}transformer." in k:
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                    
+        self._register_load_state_dict_pre_hook(_prefix_cleaner)
+
+    def tie_weights(self, **kwargs):
+        if hasattr(self, "transformer") and hasattr(self.transformer, "wte") and hasattr(self.transformer, "lm_head"):
+            self.transformer.wte.weight = self.lm_head.weight
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(self,
+                input_ids: Optional[torch.LongTensor] = None,
+                attention_mask: Optional[torch.FloatTensor] = None,
+                labels: Optional[torch.LongTensor] = None,
+                **kwargs) -> CausalLMOutputWithPast:
+        outputs = self.transformer(input_ids)
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# Register with auto-mapping
+AutoConfig.register("xpertgpt", XpertGPTConfig)
+AutoModel.register(XpertGPTConfig, XpertGPTModelWrapper)
+AutoModelForCausalLM.register(XpertGPTConfig, XpertGPTForCausalLM)

checkpoints/xpertgpt_fresh/chck_10M/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7cc187eb5c740121cab9a411116f131df4051fafaa43dfebe04293a8b817874a
+size 210797423

checkpoints/xpertgpt_fresh/chck_10M/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "[UNK]"
+}

checkpoints/xpertgpt_fresh/chck_1M/config.json

@@ -0,0 +1,17 @@
+{
+  "auto_map": {
+    "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
+    "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
+    "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
+  },
+  "vocab_size": 16384,
+  "block_size": 512,
+  "d_model": 256,
+  "hidden_size": 256,
+  "d_thin": 384,
+  "num_layers": 6,
+  "num_blocks": 4,
+  "capacity_factor": 0.5,
+  "dropout": 0.1,
+  "model_type": "xpertgpt"
+}

checkpoints/xpertgpt_fresh/chck_1M/configuration_xpertgpt.py

@@ -0,0 +1,36 @@
+from transformers import PretrainedConfig
+
+class XpertGPTConfig(PretrainedConfig):
+    model_type = "xpertgpt"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 256,
+        d_thin: int = 384,
+        num_layers: int = 6,
+        num_blocks: int = 4,
+        capacity_factor: float = 0.5,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+        self.num_hidden_layers = num_layers
+
+        super().__init__(**kwargs)

checkpoints/xpertgpt_fresh/chck_1M/modeling_xpertgpt.py

@@ -0,0 +1,456 @@
+import os
+import math
+import random
+import inspect
+from typing import Optional, Tuple, Dict, Any
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin, AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+# ─────────────────────────────────────────────────────────────
+# Configuration Classes
+# ─────────────────────────────────────────────────────────────
+
+@dataclass
+class XpertGPTModelConfig:
+    vocab_size: int = 16384
+    block_size: int = 512
+    d_model: int = 256
+    d_thin: int = 384
+    num_layers: int = 6
+    num_blocks: int = 4
+    capacity_factor: float = 0.5
+    dropout: float = 0.1
+
+class XpertGPTConfig(PretrainedConfig):
+    model_type = "xpertgpt"
+    auto_map = {
+        "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
+        "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
+        "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 256,
+        d_thin: int = 384,
+        num_layers: int = 6,
+        num_blocks: int = 4,
+        capacity_factor: float = 0.5,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+        self.num_hidden_layers = num_layers
+
+        super().__init__(**kwargs)
+
+# ─────────────────────────────────────────────────────────────
+# ROPE HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def _precompute_rope_freqs(head_dim: int, seq_len: int, device: torch.device, theta: float = 10000.0):
+    assert head_dim % 2 == 0, "head_dim must be divisible by 2 for RoPE"
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos(), emb.sin()
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    L = x.size(2)
+    cos = cos[:L, :].unsqueeze(0).unsqueeze(1) 
+    sin = sin[:L, :].unsqueeze(0).unsqueeze(1) 
+    
+    half_dim = x.size(-1) // 2
+    x1 = x[..., :half_dim]
+    x2 = x[..., half_dim:]
+    rotated_x = torch.cat((-x2, x1), dim=-1)
+    
+    return (x * cos) + (rotated_x * sin)
+
+# ─────────────────────────────────────────────────────────────
+# 1.  SLIDING WINDOW ATTENTION
+# ─────────────────────────────────────────────────────────────
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size=None):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.num_heads  = num_heads
+        self.window_size = window_size
+        self.head_dim   = dim // num_heads
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        B, L, D = x.size()
+
+        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            past_len = past_k.size(2)
+            q_cos, q_sin = _precompute_rope_freqs(self.head_dim, past_len + L, x.device)
+            q = _apply_rope(q, q_cos[past_len:, :], q_sin[past_len:, :])
+            k = _apply_rope(k, q_cos[past_len:, :], q_sin[past_len:, :])
+            
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        else:
+            cos, sin = _precompute_rope_freqs(self.head_dim, L, x.device)
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+
+        L_kv = k.size(2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        if bidirectional:
+            if self.window_size is not None:
+                past_len = L_kv - L
+                pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+                pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+                dist  = torch.abs(pos_i - pos_j)
+                win_mask = dist < self.window_size
+                scores = scores.masked_fill(
+                    ~win_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+                )
+        else:
+            past_len = L_kv - L
+            pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+            pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+            dist  = pos_i - pos_j
+
+            causal_mask = dist >= 0
+            if self.window_size is not None:
+                causal_mask = causal_mask & (dist < self.window_size)
+
+            scores = scores.masked_fill(
+                ~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+            )
+
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn, v)                                
+        out = out.transpose(1, 2).contiguous().view(B, L, D)
+        out = self.o_proj(out)
+
+        if use_cache:
+            if self.window_size is not None:
+                present_kv = (
+                    k[:, :, -self.window_size:, :],
+                    v[:, :, -self.window_size:, :]
+                )
+            else:
+                present_kv = (k, v)
+        else:
+            present_kv = None
+
+        return out, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 2.  MSIT BRANCH BLOCK (Pre-Norm)
+# ─────────────────────────────────────────────────────────────
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        # Keep parameter count comparable to regular FFN with dim * 4:
+        # 3 * dim * hidden_dim ~= 8 * dim^2  =>  hidden_dim = 8/3 * dim
+        hidden_dim = int(dim * 8 / 3)
+        hidden_dim = ((hidden_dim + 7) // 8) * 8
+        self.fc1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.fc2 = nn.Linear(dim, hidden_dim, bias=False)
+        self.fc3 = nn.Linear(hidden_dim, dim, bias=False)
+
+    def forward(self, x):
+        return self.fc3(F.silu(self.fc1(x)) * self.fc2(x))
+
+class MSITBranchBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size):
+        super().__init__()
+        self.ln1  = nn.LayerNorm(dim)
+        self.attn = SlidingWindowAttention(dim, num_heads, window_size)
+        self.ln2  = nn.LayerNorm(dim)
+        self.ffn  = SwiGLU(dim)
+        self.ln3  = nn.LayerNorm(dim)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        attn_out, present_kv = self.attn(
+            self.ln1(x), past_kv, use_cache, bidirectional
+        )
+        x = x + attn_out
+        x = x + self.ffn(self.ln2(x))
+        x = self.ln3(x)
+        return x, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 3.  MoEP-MSIT ARCHITECTURE BLOCK  (Expert Choice routing)
+# ─────────────────────────────────────────────────────────────
+
+class MoEPMSITBlock(nn.Module):
+    def __init__(self, d_model: int = 512, d_thin: int = 192, num_blocks: int = 14,
+                 capacity_factor: float = 0.5):
+        super().__init__()
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+
+        # 1. Global Block (Dense, d_model)
+        num_heads_global = max(1, d_model // 64)
+        self.global_block = MSITBranchBlock(d_model, num_heads_global, window_size=None)
+
+        # 2. Router
+        self.router_ln = nn.LayerNorm(d_model)
+        self.w_router = nn.Linear(d_model, num_blocks, bias=False)
+
+        # 3. Shrink Projection
+        self.w_down = nn.Linear(d_model, d_thin, bias=False)
+
+        # 4. Thin Parallel Blocks (d_thin)
+        self.windows = [64, 16, 8, 4] + [None] * (num_blocks - 4)
+        self.heads = [max(1, d_thin // 64)] * num_blocks  
+        
+        self.thin_blocks = nn.ModuleList([
+            MSITBranchBlock(d_thin, self.heads[i], self.windows[i])
+            for i in range(num_blocks)
+        ])
+
+        # 6. Grow Projection
+        self.w_up = nn.Linear(d_thin, d_model, bias=False)
+        self.last_topk_indices = None
+        self.ln_post_moe = nn.LayerNorm(d_model)
+
+    def forward(self, x_0: torch.Tensor, past_kvs=None, use_cache: bool = False, bidirectional: bool = False):
+        B, T, D = x_0.size()
+        n_tokens = B * T
+
+        # Step 1: Global Block
+        pkv_g = past_kvs[0] if past_kvs else None
+        x_1, nkv_g = self.global_block(x_0, pkv_g, use_cache, bidirectional)
+
+        # Step 2: Gated input stream
+        x_2 = x_1
+
+        # Router scores
+        r_logits = self.w_router(self.router_ln(x_2)).view(n_tokens, self.num_blocks)
+        r_probs = F.softmax(r_logits, dim=-1)
+
+        # Per-expert capacity: k = (n * c) / e
+        k_capacity = max(1, int(round(n_tokens * self.capacity_factor / self.num_blocks)))
+        k_capacity = min(k_capacity, n_tokens)
+
+        # Expert Choice routing: topk over the token axis for each expert
+        expert_token_scores = r_probs.transpose(0, 1) # (num_blocks, n_tokens)
+        topk_scores, topk_token_idx = torch.topk(expert_token_scores, k_capacity, dim=-1)
+        self.last_topk_indices = topk_token_idx
+
+        # Load balancing is guaranteed by construction in Expert Choice
+        layer_aux_loss = x_2.new_zeros(())
+
+        # Shrink projection
+        x_2_thin = self.w_down(x_2) # (B, T, d_thin)
+        x_2_thin_flat = x_2_thin.view(n_tokens, self.d_thin)
+
+        # Expert computations
+        new_kvs = [nkv_g]
+        expert_outputs_flat = torch.zeros(n_tokens, self.d_model, device=x_2.device, dtype=x_2.dtype)
+
+        for i, block in enumerate(self.thin_blocks):
+            sel_idx = topk_token_idx[i]
+            bucket_in = x_2_thin_flat[sel_idx].unsqueeze(0) # (1, k_capacity, d_thin)
+
+            pkv_i = past_kvs[i + 1] if past_kvs else None
+            bucket_out, nkv_i = block(bucket_in, pkv_i, use_cache, bidirectional)
+            if use_cache:
+                new_kvs.append(nkv_i)
+
+            bucket_out = bucket_out.squeeze(0) # (k_capacity, d_thin)
+            bucket_out_full = self.w_up(bucket_out) # (k_capacity, d_model)
+
+            gate = topk_scores[i].unsqueeze(-1) # (k_capacity, 1)
+            expert_outputs_flat.index_add_(0, sel_idx, bucket_out_full * gate)
+
+        x_3_full = expert_outputs_flat.view(B, T, self.d_model)
+        out = x_2 + x_3_full
+        out = self.ln_post_moe(out)
+
+        present_kvs = tuple(new_kvs) if use_cache else None
+        return out, present_kvs, layer_aux_loss
+
+# ─────────────────────────────────────────────────────────────
+# 4.  RAW XpertGPT MODEL
+# ─────────────────────────────────────────────────────────────
+
+class XpertGPTModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.wte      = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop_emb = nn.Dropout(cfg.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(
+                d_model=cfg.d_model, 
+                d_thin=cfg.d_thin, 
+                num_blocks=cfg.num_blocks, 
+                capacity_factor=cfg.capacity_factor
+            )
+            for _ in range(cfg.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(cfg.d_model)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.wte.weight = self.lm_head.weight
+
+    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None, bidirectional: bool = False):
+        x = self.drop_emb(self.wte(input_ids))
+        total_aux_loss = 0.0
+        
+        for block in self.blocks:
+            x, _, layer_aux = block(x, past_kvs=None, use_cache=False, bidirectional=bidirectional)
+            total_aux_loss += layer_aux
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            ce_loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), targets.view(-1), ignore_index=-100)
+            avg_aux_loss = total_aux_loss / self.cfg.num_layers
+            loss = ce_loss + (0.01 * avg_aux_loss)
+
+        return logits, loss
+
+# ─────────────────────────────────────────────────────────────
+# 5.  HUGGING FACE MODEL WRAPPERS
+# ─────────────────────────────────────────────────────────────
+
+class XpertGPTModelWrapper(PreTrainedModel):
+    config_class = XpertGPTConfig
+    base_model_prefix = "transformer"
+
+    def __init__(self, config: XpertGPTConfig):
+        super().__init__(config)
+        self.wte      = nn.Embedding(config.vocab_size, config.d_model)
+        self.drop_emb = nn.Dropout(config.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(config.d_model, config.d_thin, config.num_blocks, config.capacity_factor)
+            for _ in range(config.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(config.d_model)
+        self.post_init()
+
+    def forward(self, input_ids, **kwargs):
+        x = self.drop_emb(self.wte(input_ids))
+        for block in self.blocks:
+            x, _, _ = block(x, past_kvs=None, use_cache=False, bidirectional=False)
+        x = self.ln_f(x)
+        return BaseModelOutputWithPast(last_hidden_state=x)
+
+
+class XpertGPTForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = XpertGPTConfig
+    base_model_prefix = "transformer"
+    _no_split_modules = ["MoEPMSITBlock"]
+    _tied_weights_keys = {"transformer.lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config: XpertGPTConfig):
+        super().__init__(config)
+        self.transformer = XpertGPTModelWrapper(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.post_init()
+        
+        # State-dict pre-hook for backwards compatibility with checkpoint key naming
+        def _prefix_cleaner(state_dict, prefix, local_metadata, Moore, missing_keys, unexpected_keys, error_msgs):
+            keys = list(state_dict.keys())
+            for k in keys:
+                if k.startswith("transformer."):
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                elif f"{prefix}transformer." in k:
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                    
+        self._register_load_state_dict_pre_hook(_prefix_cleaner)
+
+    def tie_weights(self, **kwargs):
+        if hasattr(self, "transformer") and hasattr(self.transformer, "wte") and hasattr(self.transformer, "lm_head"):
+            self.transformer.wte.weight = self.lm_head.weight
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(self,
+                input_ids: Optional[torch.LongTensor] = None,
+                attention_mask: Optional[torch.FloatTensor] = None,
+                labels: Optional[torch.LongTensor] = None,
+                **kwargs) -> CausalLMOutputWithPast:
+        outputs = self.transformer(input_ids)
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# Register with auto-mapping
+AutoConfig.register("xpertgpt", XpertGPTConfig)
+AutoModel.register(XpertGPTConfig, XpertGPTModelWrapper)
+AutoModelForCausalLM.register(XpertGPTConfig, XpertGPTForCausalLM)

checkpoints/xpertgpt_fresh/chck_1M/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:98f17695615310b4faf9d92d2aba05d53e53225a7bc557769520d778b017414b
+size 210797423

checkpoints/xpertgpt_fresh/chck_1M/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "eos_token": "[SEP]",
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "[UNK]"
+}

checkpoints/xpertgpt_fresh/chck_20M/config.json

@@ -0,0 +1,17 @@
+{
+  "auto_map": {
+    "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
+    "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
+    "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
+  },
+  "vocab_size": 16384,
+  "block_size": 512,
+  "d_model": 256,
+  "hidden_size": 256,
+  "d_thin": 384,
+  "num_layers": 6,
+  "num_blocks": 4,
+  "capacity_factor": 0.5,
+  "dropout": 0.1,
+  "model_type": "xpertgpt"
+}

checkpoints/xpertgpt_fresh/chck_20M/configuration_xpertgpt.py

@@ -0,0 +1,36 @@
+from transformers import PretrainedConfig
+
+class XpertGPTConfig(PretrainedConfig):
+    model_type = "xpertgpt"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 256,
+        d_thin: int = 384,
+        num_layers: int = 6,
+        num_blocks: int = 4,
+        capacity_factor: float = 0.5,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+        self.num_hidden_layers = num_layers
+
+        super().__init__(**kwargs)

checkpoints/xpertgpt_fresh/chck_20M/modeling_xpertgpt.py

@@ -0,0 +1,456 @@
+import os
+import math
+import random
+import inspect
+from typing import Optional, Tuple, Dict, Any
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin, AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+# ─────────────────────────────────────────────────────────────
+# Configuration Classes
+# ─────────────────────────────────────────────────────────────
+
+@dataclass
+class XpertGPTModelConfig:
+    vocab_size: int = 16384
+    block_size: int = 512
+    d_model: int = 256
+    d_thin: int = 384
+    num_layers: int = 6
+    num_blocks: int = 4
+    capacity_factor: float = 0.5
+    dropout: float = 0.1
+
+class XpertGPTConfig(PretrainedConfig):
+    model_type = "xpertgpt"
+    auto_map = {
+        "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
+        "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
+        "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 256,
+        d_thin: int = 384,
+        num_layers: int = 6,
+        num_blocks: int = 4,
+        capacity_factor: float = 0.5,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+        self.num_hidden_layers = num_layers
+
+        super().__init__(**kwargs)
+
+# ─────────────────────────────────────────────────────────────
+# ROPE HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def _precompute_rope_freqs(head_dim: int, seq_len: int, device: torch.device, theta: float = 10000.0):
+    assert head_dim % 2 == 0, "head_dim must be divisible by 2 for RoPE"
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos(), emb.sin()
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    L = x.size(2)
+    cos = cos[:L, :].unsqueeze(0).unsqueeze(1) 
+    sin = sin[:L, :].unsqueeze(0).unsqueeze(1) 
+    
+    half_dim = x.size(-1) // 2
+    x1 = x[..., :half_dim]
+    x2 = x[..., half_dim:]
+    rotated_x = torch.cat((-x2, x1), dim=-1)
+    
+    return (x * cos) + (rotated_x * sin)
+
+# ─────────────────────────────────────────────────────────────
+# 1.  SLIDING WINDOW ATTENTION
+# ─────────────────────────────────────────────────────────────
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size=None):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.num_heads  = num_heads
+        self.window_size = window_size
+        self.head_dim   = dim // num_heads
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        B, L, D = x.size()
+
+        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            past_len = past_k.size(2)
+            q_cos, q_sin = _precompute_rope_freqs(self.head_dim, past_len + L, x.device)
+            q = _apply_rope(q, q_cos[past_len:, :], q_sin[past_len:, :])
+            k = _apply_rope(k, q_cos[past_len:, :], q_sin[past_len:, :])
+            
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        else:
+            cos, sin = _precompute_rope_freqs(self.head_dim, L, x.device)
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+
+        L_kv = k.size(2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        if bidirectional:
+            if self.window_size is not None:
+                past_len = L_kv - L
+                pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+                pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+                dist  = torch.abs(pos_i - pos_j)
+                win_mask = dist < self.window_size
+                scores = scores.masked_fill(
+                    ~win_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+                )
+        else:
+            past_len = L_kv - L
+            pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+            pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+            dist  = pos_i - pos_j
+
+            causal_mask = dist >= 0
+            if self.window_size is not None:
+                causal_mask = causal_mask & (dist < self.window_size)
+
+            scores = scores.masked_fill(
+                ~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+            )
+
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn, v)                                
+        out = out.transpose(1, 2).contiguous().view(B, L, D)
+        out = self.o_proj(out)
+
+        if use_cache:
+            if self.window_size is not None:
+                present_kv = (
+                    k[:, :, -self.window_size:, :],
+                    v[:, :, -self.window_size:, :]
+                )
+            else:
+                present_kv = (k, v)
+        else:
+            present_kv = None
+
+        return out, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 2.  MSIT BRANCH BLOCK (Pre-Norm)
+# ─────────────────────────────────────────────────────────────
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        # Keep parameter count comparable to regular FFN with dim * 4:
+        # 3 * dim * hidden_dim ~= 8 * dim^2  =>  hidden_dim = 8/3 * dim
+        hidden_dim = int(dim * 8 / 3)
+        hidden_dim = ((hidden_dim + 7) // 8) * 8
+        self.fc1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.fc2 = nn.Linear(dim, hidden_dim, bias=False)
+        self.fc3 = nn.Linear(hidden_dim, dim, bias=False)
+
+    def forward(self, x):
+        return self.fc3(F.silu(self.fc1(x)) * self.fc2(x))
+
+class MSITBranchBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size):
+        super().__init__()
+        self.ln1  = nn.LayerNorm(dim)
+        self.attn = SlidingWindowAttention(dim, num_heads, window_size)
+        self.ln2  = nn.LayerNorm(dim)
+        self.ffn  = SwiGLU(dim)
+        self.ln3  = nn.LayerNorm(dim)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        attn_out, present_kv = self.attn(
+            self.ln1(x), past_kv, use_cache, bidirectional
+        )
+        x = x + attn_out
+        x = x + self.ffn(self.ln2(x))
+        x = self.ln3(x)
+        return x, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 3.  MoEP-MSIT ARCHITECTURE BLOCK  (Expert Choice routing)
+# ─────────────────────────────────────────────────────────────
+
+class MoEPMSITBlock(nn.Module):
+    def __init__(self, d_model: int = 512, d_thin: int = 192, num_blocks: int = 14,
+                 capacity_factor: float = 0.5):
+        super().__init__()
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+
+        # 1. Global Block (Dense, d_model)
+        num_heads_global = max(1, d_model // 64)
+        self.global_block = MSITBranchBlock(d_model, num_heads_global, window_size=None)
+
+        # 2. Router
+        self.router_ln = nn.LayerNorm(d_model)
+        self.w_router = nn.Linear(d_model, num_blocks, bias=False)
+
+        # 3. Shrink Projection
+        self.w_down = nn.Linear(d_model, d_thin, bias=False)
+
+        # 4. Thin Parallel Blocks (d_thin)
+        self.windows = [64, 16, 8, 4] + [None] * (num_blocks - 4)
+        self.heads = [max(1, d_thin // 64)] * num_blocks  
+        
+        self.thin_blocks = nn.ModuleList([
+            MSITBranchBlock(d_thin, self.heads[i], self.windows[i])
+            for i in range(num_blocks)
+        ])
+
+        # 6. Grow Projection
+        self.w_up = nn.Linear(d_thin, d_model, bias=False)
+        self.last_topk_indices = None
+        self.ln_post_moe = nn.LayerNorm(d_model)
+
+    def forward(self, x_0: torch.Tensor, past_kvs=None, use_cache: bool = False, bidirectional: bool = False):
+        B, T, D = x_0.size()
+        n_tokens = B * T
+
+        # Step 1: Global Block
+        pkv_g = past_kvs[0] if past_kvs else None
+        x_1, nkv_g = self.global_block(x_0, pkv_g, use_cache, bidirectional)
+
+        # Step 2: Gated input stream
+        x_2 = x_1
+
+        # Router scores
+        r_logits = self.w_router(self.router_ln(x_2)).view(n_tokens, self.num_blocks)
+        r_probs = F.softmax(r_logits, dim=-1)
+
+        # Per-expert capacity: k = (n * c) / e
+        k_capacity = max(1, int(round(n_tokens * self.capacity_factor / self.num_blocks)))
+        k_capacity = min(k_capacity, n_tokens)
+
+        # Expert Choice routing: topk over the token axis for each expert
+        expert_token_scores = r_probs.transpose(0, 1) # (num_blocks, n_tokens)
+        topk_scores, topk_token_idx = torch.topk(expert_token_scores, k_capacity, dim=-1)
+        self.last_topk_indices = topk_token_idx
+
+        # Load balancing is guaranteed by construction in Expert Choice
+        layer_aux_loss = x_2.new_zeros(())
+
+        # Shrink projection
+        x_2_thin = self.w_down(x_2) # (B, T, d_thin)
+        x_2_thin_flat = x_2_thin.view(n_tokens, self.d_thin)
+
+        # Expert computations
+        new_kvs = [nkv_g]
+        expert_outputs_flat = torch.zeros(n_tokens, self.d_model, device=x_2.device, dtype=x_2.dtype)
+
+        for i, block in enumerate(self.thin_blocks):
+            sel_idx = topk_token_idx[i]
+            bucket_in = x_2_thin_flat[sel_idx].unsqueeze(0) # (1, k_capacity, d_thin)
+
+            pkv_i = past_kvs[i + 1] if past_kvs else None
+            bucket_out, nkv_i = block(bucket_in, pkv_i, use_cache, bidirectional)
+            if use_cache:
+                new_kvs.append(nkv_i)
+
+            bucket_out = bucket_out.squeeze(0) # (k_capacity, d_thin)
+            bucket_out_full = self.w_up(bucket_out) # (k_capacity, d_model)
+
+            gate = topk_scores[i].unsqueeze(-1) # (k_capacity, 1)
+            expert_outputs_flat.index_add_(0, sel_idx, bucket_out_full * gate)
+
+        x_3_full = expert_outputs_flat.view(B, T, self.d_model)
+        out = x_2 + x_3_full
+        out = self.ln_post_moe(out)
+
+        present_kvs = tuple(new_kvs) if use_cache else None
+        return out, present_kvs, layer_aux_loss
+
+# ─────────────────────────────────────────────────────────────
+# 4.  RAW XpertGPT MODEL
+# ─────────────────────────────────────────────────────────────
+
+class XpertGPTModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.wte      = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop_emb = nn.Dropout(cfg.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(
+                d_model=cfg.d_model, 
+                d_thin=cfg.d_thin, 
+                num_blocks=cfg.num_blocks, 
+                capacity_factor=cfg.capacity_factor
+            )
+            for _ in range(cfg.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(cfg.d_model)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.wte.weight = self.lm_head.weight
+
+    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None, bidirectional: bool = False):
+        x = self.drop_emb(self.wte(input_ids))
+        total_aux_loss = 0.0
+        
+        for block in self.blocks:
+            x, _, layer_aux = block(x, past_kvs=None, use_cache=False, bidirectional=bidirectional)
+            total_aux_loss += layer_aux
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            ce_loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), targets.view(-1), ignore_index=-100)
+            avg_aux_loss = total_aux_loss / self.cfg.num_layers
+            loss = ce_loss + (0.01 * avg_aux_loss)
+
+        return logits, loss
+
+# ─────────────────────────────────────────────────────────────
+# 5.  HUGGING FACE MODEL WRAPPERS
+# ─────────────────────────────────────────────────────────────
+
+class XpertGPTModelWrapper(PreTrainedModel):
+    config_class = XpertGPTConfig
+    base_model_prefix = "transformer"
+
+    def __init__(self, config: XpertGPTConfig):
+        super().__init__(config)
+        self.wte      = nn.Embedding(config.vocab_size, config.d_model)
+        self.drop_emb = nn.Dropout(config.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(config.d_model, config.d_thin, config.num_blocks, config.capacity_factor)
+            for _ in range(config.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(config.d_model)
+        self.post_init()
+
+    def forward(self, input_ids, **kwargs):
+        x = self.drop_emb(self.wte(input_ids))
+        for block in self.blocks:
+            x, _, _ = block(x, past_kvs=None, use_cache=False, bidirectional=False)
+        x = self.ln_f(x)
+        return BaseModelOutputWithPast(last_hidden_state=x)
+
+
+class XpertGPTForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = XpertGPTConfig
+    base_model_prefix = "transformer"
+    _no_split_modules = ["MoEPMSITBlock"]
+    _tied_weights_keys = {"transformer.lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config: XpertGPTConfig):
+        super().__init__(config)
+        self.transformer = XpertGPTModelWrapper(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.post_init()
+        
+        # State-dict pre-hook for backwards compatibility with checkpoint key naming
+        def _prefix_cleaner(state_dict, prefix, local_metadata, Moore, missing_keys, unexpected_keys, error_msgs):
+            keys = list(state_dict.keys())
+            for k in keys:
+                if k.startswith("transformer."):
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                elif f"{prefix}transformer." in k:
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                    
+        self._register_load_state_dict_pre_hook(_prefix_cleaner)
+
+    def tie_weights(self, **kwargs):
+        if hasattr(self, "transformer") and hasattr(self.transformer, "wte") and hasattr(self.transformer, "lm_head"):
+            self.transformer.wte.weight = self.lm_head.weight
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(self,
+                input_ids: Optional[torch.LongTensor] = None,
+                attention_mask: Optional[torch.FloatTensor] = None,
+                labels: Optional[torch.LongTensor] = None,
+                **kwargs) -> CausalLMOutputWithPast:
+        outputs = self.transformer(input_ids)
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# Register with auto-mapping
+AutoConfig.register("xpertgpt", XpertGPTConfig)
+AutoModel.register(XpertGPTConfig, XpertGPTModelWrapper)
+AutoModelForCausalLM.register(XpertGPTConfig, XpertGPTForCausalLM)

checkpoints/xpertgpt_fresh/chck_20M/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e9078c28108f48e47eb7600fd1dfbf7ef616cf8def294f34e3d82d22b1f0838
+size 210797423