Repository: LAION-AI/CLAP (language: python)

Issue:
Code merge
Contains all the code update for ICASSP 2023 paper


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 45acf18..6f01146 100644
--- a/.gitignore
+++ b/.gitignore
@@ -154,3 +154,11 @@ src/sizes.json
 *.pt
 dev.ipynb
 src/json_files/*
+.vs/*
+run.sh
+Untitled.ipynb
+dev.ipynb
+src/dev.py
+ESC.csv
+.gitignore
+UrbanSound8K.csv
diff --git a/ESC50_class_labels_indices.json b/ESC50_class_labels_indices.json
new file mode 100644
index 0000000..d860536
--- /dev/null
+++ b/ESC50_class_labels_indices.json
@@ -0,0 +1,1 @@
+{"dog": 0, "rooster": 1, "pig": 2, "cow": 3, "frog": 4, "cat": 5, "hen": 6, "insects": 7, "sheep": 8, "crow": 9, "rain": 10, "sea_waves": 11, "crackling_fire": 12, "crickets": 13, "chirping_birds": 14, "water_drops": 15, "wind": 16, "pouring_water": 17, "toilet_flush": 18, "thunderstorm": 19, "crying_baby": 20, "sneezing": 21, "clapping": 22, "breathing": 23, "coughing": 24, "footsteps": 25, "laughing": 26, "brushing_teeth": 27, "snoring": 28, "drinking_sipping": 29, "door_wood_knock": 30, "mouse_click": 31, "keyboard_typing": 32, "door_wood_creaks": 33, "can_opening": 34, "washing_machine": 35, "vacuum_cleaner": 36, "clock_alarm": 37, "clock_tick": 38, "glass_breaking": 39, "helicopter": 40, "chainsaw": 41, "siren": 42, "car_horn": 43, "engine": 44, "train": 45, "church_bells": 46, "airplane": 47, "fireworks": 48, "hand_saw": 49}
\ No newline at end of file
diff --git a/README.md b/README.md
index 251d945..80158ee 100644
--- a/README.md
+++ b/README.md
@@ -86,9 +86,14 @@ torchrun --nproc_per_node 8  -m training.main \
     --resample-method="None" \
     --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" "audioset" "free_to_use_sounds" "paramount_motion" "sonniss_game_effects" "wesoundeffects" \
     --datasetinfos "train" "unbalanced_train" "balanced_train" \
-    --datasetpath <webdataset_tar_path>
+    --datasetpath <webdataset_tar_path> \
+    --data-filling repeat \
+    --kappa 1.0 \
 ```
 
+**kappa** is the temperature parameter for weighted contrastive CLIPLoss. Set **data-filling** to **repeatpad** means when you have a 3s audio, and the max_len is 10s, it will repeat 3 times to 9s and then pad to 10s with zeros. Set **data-filling** to **pad** means when you have a 3s audio, and the max_len is 10s, it will pad to 10s. Set **data-filling** to **repeat** means when you have a 3s audio, and the max_len is 10s, it will repeat to 12s and clip to 10s. 
+
+
 ## How to get audio feature?
 Please refer to the ``evaluation()`` function in the ``train.py`` and `audio_infer()` function in `model.py`.
 
@@ -98,3 +103,16 @@ Please refer to the ``get_metrics()`` function in the ``train.py``.
 ## Link to Pretrained Models
 TODO
 
+## Test if tar is invalid
+``python test_tars.py --tar-path "pipe:aws s3 --cli-connect-timeout 0 cp s3://s-laion-audio/webdataset_tar/audioset/unbalanced_train" --start 0 --end 100 --exclude 11 21 --batch-size 1 --order``
+
+This means test tars from ``pipe:aws s3 --cli-connect-timeout 0 cp s3://s-laion-audio/webdataset_tar/audioset/unbalanced_train/0.tar`` to ``pipe:aws s3 --cli-connect-timeout 0 cp s3://s-laion-audio/webdataset_tar/audioset/unbalanced_train/100.tar`` but exclude ``pipe:aws s3 --cli-connect-timeout 0 cp s3://s-laion-audio/webdataset_tar/audioset/unbalanced_train/11.tar``and ``pipe:aws s3 --cli-connect-timeout 0 cp s3://s-laion-audio/webdataset_tar/audioset/unbalanced_train/21.tar``. The iterative order if not random by specifying `--order`.
+
+## Test if checkpoints are freezed
+
+Use ``keys_in_state_dict()`` in ``check_ckpt.py`` to check the keys in the state_dict. Use ``check_ckpt_diff`` to check if 2 checkpoints has the same weights. e.g.
+
+``check_ckpt_diff("/fsx/clap_logs/2022_09_11-19_37_08-model_PANN-14-lr_0.001-b_160-j_4-p_fp32/checkpoints/epoch_10.pt", "/fsx/clap_logs/2022_09_11-19_37_08-model_PANN-14-lr_0.001-b_160-j_4-p_fp32/checkpoints/epoch_100.pt", "text_branch.resblocks")``
+
+
+This means check the difference between ``/fsx/clap_logs/2022_09_11-19_37_08-model_PANN-14-lr_0.001-b_160-j_4-p_fp32/checkpoints/epoch_10.pt`` and ``/fsx/clap_logs/2022_09_11-19_37_08-model_PANN-14-lr_0.001-b_160-j_4-p_fp32/checkpoints/epoch_100.pt``, easpecally with keys include ``text_branch.resblocks``.
diff --git a/UrbanSound8K_class_labels_indices.json b/UrbanSound8K_class_labels_indices.json
new file mode 100644
index 0000000..dc94116
--- /dev/null
+++ b/UrbanSound8K_class_labels_indices.json
@@ -0,0 +1,1 @@
+{"air_conditioner": 0, "car_horn": 1, "children_playing": 2, "dog_bark": 3, "drilling": 4, "engine_idling": 5, "gun_shot": 6, "jackhammer": 7, "siren": 8, "street_music": 9}
\ No newline at end of file
diff --git a/audioset_class_labels_indices.json b/audioset_class_labels_indices.json
new file mode 100644
index 0000000..0f6780b
--- /dev/null
+++ b/audioset_class_labels_indices.json
@@ -0,0 +1,1 @@
+{"Speech": 0, "Male speech, man speaking": 1, "Female speech, woman speaking": 2, "Child speech, kid speaking": 3, "Conversation": 4, "Narration, monologue": 5, "Babbling": 6, "Speech synthesizer": 7, "Shout": 8, "Bellow": 9, "Whoop": 10, "Yell": 11, "Battle cry": 12, "Children shouting": 13, "Screaming": 14, "Whispering": 15, "Laughter": 16, "Baby laughter": 17, "Giggle": 18, "Snicker": 19, "Belly laugh": 20, "Chuckle, chortle": 21, "Crying, sobbing": 22, "Baby cry, infant cry": 23, "Whimper": 24, "Wail, moan": 25, "Sigh": 26, "Singing": 27, "Choir": 28, "Yodeling": 29, "Chant": 30, "Mantra": 31, "Male singing": 32, "Female singing": 33, "Child singing": 34, "Synthetic singing": 35, "Rapping": 36, "Humming": 37, "Groan": 38, "Grunt": 39, "Whistling": 40, "Breathing": 41, "Wheeze": 42, "Snoring": 43, "Gasp": 44, "Pant": 45, "Snort": 46, "Cough": 47, "Throat clearing": 48, "Sneeze": 49, "Sniff": 50, "Run": 51, "Shuffle": 52, "Walk, footsteps": 53, "Chewing, mastication": 54, "Biting": 55, "Gargling": 56, "Stomach rumble": 57, "Burping, eructation": 58, "Hiccup": 59, "Fart": 60, "Hands": 61, "Finger snapping": 62, "Clapping": 63, "Heart sounds, heartbeat": 64, "Heart murmur": 65, "Cheering": 66, "Applause": 67, "Chatter": 68, "Crowd": 69, "Hubbub, speech noise, speech babble": 70, "Children playing": 71, "Animal": 72, "Domestic animals, pets": 73, "Dog": 74, "Bark": 75, "Yip": 76, "Howl": 77, "Bow-wow": 78, "Growling": 79, "Whimper (dog)": 80, "Cat": 81, "Purr": 82, "Meow": 83, "Hiss": 84, "Caterwaul": 85, "Livestock, farm animals, working animals": 86, "Horse": 87, "Clip-clop": 88, "Neigh, whinny": 89, "Cattle, bovinae": 90, "Moo": 91, "Cowbell": 92, "Pig": 93, "Oink": 94, "Goat": 95, "Bleat": 96, "Sheep": 97, "Fowl": 98, "Chicken, rooster": 99, "Cluck": 100, "Crowing, cock-a-doodle-doo": 101, "Turkey": 102, "Gobble": 103, "Duck": 104, "Quack": 105, "Goose": 106, "Honk": 107, "Wild animals": 108, "Roaring cats (lions, tigers)": 109, "Roar": 110, "Bird": 111, "Bird vocalization, bird call, bird song": 112, "Chirp, tweet": 113, "Squawk": 114, "Pigeon, dove": 115, "Coo": 116, "Crow": 117, "Caw": 118, "Owl": 119, "Hoot": 120, "Bird flight, flapping wings": 121, "Canidae, dogs, wolves": 122, "Rodents, rats, mice": 123, "Mouse": 124, "Patter": 125, "Insect": 126, "Cricket": 127, "Mosquito": 128, "Fly, housefly": 129, "Buzz": 130, "Bee, wasp, etc.": 131, "Frog": 132, "Croak": 133, "Snake": 134, "Rattle": 135, "Whale vocalization": 136, "Music": 137, "Musical instrument": 138, "Plucked string instrument": 139, "Guitar": 140, "Electric guitar": 141, "Bass guitar": 142, "Acoustic guitar": 143, "Steel guitar, slide guitar": 144, "Tapping (guitar technique)": 145, "Strum": 146, "Banjo": 147, "Sitar": 148, "Mandolin": 149, "Zither": 150, "Ukulele": 151, "Keyboard (musical)": 152, "Piano": 153, "Electric piano": 154, "Organ": 155, "Electronic organ": 156, "Hammond organ": 157, "Synthesizer": 158, "Sampler": 159, "Harpsichord": 160, "Percussion": 161, "Drum kit": 162, "Drum machine": 163, "Drum": 164, "Snare drum": 165, "Rimshot": 166, "Drum roll": 167, "Bass drum": 168, "Timpani": 169, "Tabla": 170, "Cymbal": 171, "Hi-hat": 172, "Wood block": 173, "Tambourine": 174, "Rattle (instrument)": 175, "Maraca": 176, "Gong": 177, "Tubular bells": 178, "Mallet percussion": 179, "Marimba, xylophone": 180, "Glockenspiel": 181, "Vibraphone": 182, "Steelpan": 183, "Orchestra": 184, "Brass instrument": 185, "French horn": 186, "Trumpet": 187, "Trombone": 188, "Bowed string instrument": 189, "String section": 190, "Violin, fiddle": 191, "Pizzicato": 192, "Cello": 193, "Double bass": 194, "Wind instrument, woodwind instrument": 195, "Flute": 196, "Saxophone": 197, "Clarinet": 198, "Harp": 199, "Bell": 200, "Church bell": 201, "Jingle bell": 202, "Bicycle bell": 203, "Tuning fork": 204, "Chime": 205, "Wind chime": 206, "Change ringing (campanology)": 207, "Harmonica": 208, "Accordion": 209, "Bagpipes": 210, "Didgeridoo": 211, "Shofar": 212, "Theremin": 213, "Singing bowl": 214, "Scratching (performance technique)": 215, "Pop music": 216, "Hip hop music": 217, "Beatboxing": 218, "Rock music": 219, "Heavy metal": 220, "Punk rock": 221, "Grunge": 222, "Progressive rock": 223, "Rock and roll": 224, "Psychedelic rock": 225, "Rhythm and blues": 226, "Soul music": 227, "Reggae": 228, "Country": 229, "Swing music": 230, "Bluegrass": 231, "Funk": 232, "Folk music": 233, "Middle Eastern music": 234, "Jazz": 235, "Disco": 236, "Classical music": 237, "Opera": 238, "Electronic music": 239, "House music": 240, "Techno": 241, "Dubstep": 242, "Drum and bass": 243, "Electronica": 244, "Electronic dance music": 245, "Ambient music": 246, "Trance music": 247, "Music of Latin America": 248, "Salsa music": 249, "Flamenco": 250, "Blues": 251, "Music for children": 252, "New-age music": 253, "Vocal music": 254, "A capella": 255, "Music of Africa": 256, "Afrobeat": 257, "Christian music": 258, "Gospel music": 259, "Music of Asia": 260, "Carnatic music": 261, "Music of Bollywood": 262, "Ska": 263, "Traditional music": 264, "Independent music": 265, "Song": 266, "Background music": 267, "Theme music": 268, "Jingle (music)": 269, "Soundtrack music": 270, "Lullaby": 271, "Video game music": 272, "Christmas music": 273, "Dance music": 274, "Wedding music": 275, "Happy music": 276, "Funny music": 277, "Sad music": 278, "Tender music": 279, "Exciting music": 280, "Angry music": 281, "Scary music": 282, "Wind": 283, "Rustling leaves": 284, "Wind noise (microphone)": 285, "Thunderstorm": 286, "Thunder": 287, "Water": 288, "Rain": 289, "Raindrop": 290, "Rain on surface": 291, "Stream": 292, "Waterfall": 293, "Ocean": 294, "Waves, surf": 295, "Steam": 296, "Gurgling": 297, "Fire": 298, "Crackle": 299, "Vehicle": 300, "Boat, Water vehicle": 301, "Sailboat, sailing ship": 302, "Rowboat, canoe, kayak": 303, "Motorboat, speedboat": 304, "Ship": 305, "Motor vehicle (road)": 306, "Car": 307, "Vehicle horn, car horn, honking": 308, "Toot": 309, "Car alarm": 310, "Power windows, electric windows": 311, "Skidding": 312, "Tire squeal": 313, "Car passing by": 314, "Race car, auto racing": 315, "Truck": 316, "Air brake": 317, "Air horn, truck horn": 318, "Reversing beeps": 319, "Ice cream truck, ice cream van": 320, "Bus": 321, "Emergency vehicle": 322, "Police car (siren)": 323, "Ambulance (siren)": 324, "Fire engine, fire truck (siren)": 325, "Motorcycle": 326, "Traffic noise, roadway noise": 327, "Rail transport": 328, "Train": 329, "Train whistle": 330, "Train horn": 331, "Railroad car, train wagon": 332, "Train wheels squealing": 333, "Subway, metro, underground": 334, "Aircraft": 335, "Aircraft engine": 336, "Jet engine": 337, "Propeller, airscrew": 338, "Helicopter": 339, "Fixed-wing aircraft, airplane": 340, "Bicycle": 341, "Skateboard": 342, "Engine": 343, "Light engine (high frequency)": 344, "Dental drill, dentist's drill": 345, "Lawn mower": 346, "Chainsaw": 347, "Medium engine (mid frequency)": 348, "Heavy engine (low frequency)": 349, "Engine knocking": 350, "Engine starting": 351, "Idling": 352, "Accelerating, revving, vroom": 353, "Door": 354, "Doorbell": 355, "Ding-dong": 356, "Sliding door": 357, "Slam": 358, "Knock": 359, "Tap": 360, "Squeak": 361, "Cupboard open or close": 362, "Drawer open or close": 363, "Dishes, pots, and pans": 364, "Cutlery, silverware": 365, "Chopping (food)": 366, "Frying (food)": 367, "Microwave oven": 368, "Blender": 369, "Water tap, faucet": 370, "Sink (filling or washing)": 371, "Bathtub (filling or washing)": 372, "Hair dryer": 373, "Toilet flush": 374, "Toothbrush": 375, "Electric toothbrush": 376, "Vacuum cleaner": 377, "Zipper (clothing)": 378, "Keys jangling": 379, "Coin (dropping)": 380, "Scissors": 381, "Electric shaver, electric razor": 382, "Shuffling cards": 383, "Typing": 384, "Typewriter": 385, "Computer keyboard": 386, "Writing": 387, "Alarm": 388, "Telephone": 389, "Telephone bell ringing": 390, "Ringtone": 391, "Telephone dialing, DTMF": 392, "Dial tone": 393, "Busy signal": 394, "Alarm clock": 395, "Siren": 396, "Civil defense siren": 397, "Buzzer": 398, "Smoke detector, smoke alarm": 399, "Fire alarm": 400, "Foghorn": 401, "Whistle": 402, "Steam whistle": 403, "Mechanisms": 404, "Ratchet, pawl": 405, "Clock": 406, "Tick": 407, "Tick-tock": 408, "Gears": 409, "Pulleys": 410, "Sewing machine": 411, "Mechanical fan": 412, "Air conditioning": 413, "Cash register": 414, "Printer": 415, "Camera": 416, "Single-lens reflex camera": 417, "Tools": 418, "Hammer": 419, "Jackhammer": 420, "Sawing": 421, "Filing (rasp)": 422, "Sanding": 423, "Power tool": 424, "Drill": 425, "Explosion": 426, "Gunshot, gunfire": 427, "Machine gun": 428, "Fusillade": 429, "Artillery fire": 430, "Cap gun": 431, "Fireworks": 432, "Firecracker": 433, "Burst, pop": 434, "Eruption": 435, "Boom": 436, "Wood": 437, "Chop": 438, "Splinter": 439, "Crack": 440, "Glass": 441, "Chink, clink": 442, "Shatter": 443, "Liquid": 444, "Splash, splatter": 445, "Slosh": 446, "Squish": 447, "Drip": 448, "Pour": 449, "Trickle, dribble": 450, "Gush": 451, "Fill (with liquid)": 452, "Spray": 453, "Pump (liquid)": 454, "Stir": 455, "Boiling": 456, "Sonar": 457, "Arrow": 458, "Whoosh, swoosh, swish": 459, "Thump, thud": 460, "Thunk": 461, "Electronic tuner": 462, "Effects unit": 463, "Chorus effect": 464, "Basketball bounce": 465, "Bang": 466, "Slap, smack": 467, "Whack, thwack": 468, "Smash, crash": 469, "Breaking": 470, "Bouncing": 471, "Whip": 472, "Flap": 473, "Scratch": 474, "Scrape": 475, "Rub": 476, "Roll": 477, "Crushing": 478, "Crumpling, crinkling": 479, "Tearing": 480, "Beep, bleep": 481, "Ping": 482, "Ding": 483, "Clang": 484, "Squeal": 485, "Creak": 486, "Rustle": 487, "Whir": 488, "Clatter": 489, "Sizzle": 490, "Clicking": 491, "Clickety-clack": 492, "Rumble": 493, "Plop": 494, "Jingle, tinkle": 495, "Hum": 496, "Zing": 497, "Boing": 498, "Crunch": 499, "Silence": 500, "Sine wave": 501, "Harmonic": 502, "Chirp tone": 503, "Sound effect": 504, "Pulse": 505, "Inside, small room": 506, "Inside, large room or hall": 507, "Inside, public space": 508, "Outside, urban or manmade": 509, "Outside, rural or natural": 510, "Reverberation": 511, "Echo": 512, "Noise": 513, "Environmental noise": 514, "Static": 515, "Mains hum": 516, "Distortion": 517, "Sidetone": 518, "Cacophony": 519, "White noise": 520, "Pink noise": 521, "Throbbing": 522, "Vibration": 523, "Television": 524, "Radio": 525, "Field recording": 526}
\ No newline at end of file
diff --git a/class_labels/ESC50_class_labels_indices.json b/class_labels/ESC50_class_labels_indices.json
new file mode 100644
index 0000000..d860536
--- /dev/null
+++ b/class_labels/ESC50_class_labels_indices.json
@@ -0,0 +1,1 @@
+{"dog": 0, "rooster": 1, "pig": 2, "cow": 3, "frog": 4, "cat": 5, "hen": 6, "insects": 7, "sheep": 8, "crow": 9, "rain": 10, "sea_waves": 11, "crackling_fire": 12, "crickets": 13, "chirping_birds": 14, "water_drops": 15, "wind": 16, "pouring_water": 17, "toilet_flush": 18, "thunderstorm": 19, "crying_baby": 20, "sneezing": 21, "clapping": 22, "breathing": 23, "coughing": 24, "footsteps": 25, "laughing": 26, "brushing_teeth": 27, "snoring": 28, "drinking_sipping": 29, "door_wood_knock": 30, "mouse_click": 31, "keyboard_typing": 32, "door_wood_creaks": 33, "can_opening": 34, "washing_machine": 35, "vacuum_cleaner": 36, "clock_alarm": 37, "clock_tick": 38, "glass_breaking": 39, "helicopter": 40, "chainsaw": 41, "siren": 42, "car_horn": 43, "engine": 44, "train": 45, "church_bells": 46, "airplane": 47, "fireworks": 48, "hand_saw": 49}
\ No newline at end of file
diff --git a/class_labels/ESC50_class_labels_indices_space.json b/class_labels/ESC50_class_labels_indices_space.json
new file mode 100644
index 0000000..1293de7
--- /dev/null
+++ b/class_labels/ESC50_class_labels_indices_space.json
@@ -0,0 +1,1 @@
+{"dog": 0, "rooster": 1, "pig": 2, "cow": 3, "frog": 4, "cat": 5, "hen": 6, "insects": 7, "sheep": 8, "crow": 9, "rain": 10, "sea waves": 11, "crackling fire": 12, "crickets": 13, "chirping birds": 14, "water drops": 15, "wind": 16, "pouring water": 17, "toilet flush": 18, "thunderstorm": 19, "crying baby": 20, "sneezing": 21, "clapping": 22, "breathing": 23, "coughing": 24, "footsteps": 25, "laughing": 26, "brushing teeth": 27, "snoring": 28, "drinking sipping": 29, "door wood knock": 30, "mouse click": 31, "keyboard typing": 32, "door wood creaks": 33, "can opening": 34, "washing machine": 35, "vacuum cleaner": 36, "clock alarm": 37, "clock tick": 38, "glass breaking": 39, "helicopter": 40, "chainsaw": 41, "siren": 42, "car horn": 43, "engine": 44, "train": 45, "church bells": 46, "airplane": 47, "fireworks": 48, "hand saw": 49}
\ No newline at end of file
diff --git a/class_labels/FSD50k_class_labels_indices.json b/class_labels/FSD50k_class_labels_indices.json
new file mode 100644
index 0000000..23d98f2
--- /dev/null
+++ b/class_labels/FSD50k_class_labels_indices.json
@@ -0,0 +1,1 @@
+{"Whispering": 0, "Gunshot, gunfire": 1, "Pour": 2, "Wind chime": 3, "Livestock, farm animals, working animals": 4, "Crackle": 5, "Waves, surf": 6, "Chicken, rooster": 7, "Chatter": 8, "Keyboard (musical)": 9, "Bark": 10, "Rail transport": 11, "Gong": 12, "Shatter": 13, "Ratchet, pawl": 14, "Clapping": 15, "Mallet percussion": 16, "Whoosh, swoosh, swish": 17, "Speech synthesizer": 18, "Respiratory sounds": 19, "Sliding door": 20, "Boat, Water vehicle": 21, "Boiling": 22, "Human voice": 23, "Drip": 24, "Thunderstorm": 25, "Male singing": 26, "Sneeze": 27, "Hi-hat": 28, "Guitar": 29, "Crying, sobbing": 30, "Speech": 31, "Slam": 32, "Crack": 33, "Yell": 34, "Drawer open or close": 35, "Run": 36, "Cheering": 37, "Splash, splatter": 38, "Tabla": 39, "Sigh": 40, "Packing tape, duct tape": 41, "Raindrop": 42, "Cymbal": 43, "Fill (with liquid)": 44, "Harp": 45, "Squeak": 46, "Zipper (clothing)": 47, "Tearing": 48, "Alarm": 49, "Skateboard": 50, "Wind instrument, woodwind instrument": 51, "Chink, clink": 52, "Wind": 53, "Ringtone": 54, "Microwave oven": 55, "Power tool": 56, "Dishes, pots, and pans": 57, "Musical instrument": 58, "Door": 59, "Domestic sounds, home sounds": 60, "Subway, metro, underground": 61, "Glockenspiel": 62, "Female speech, woman speaking": 63, "Coin (dropping)": 64, "Mechanical fan": 65, "Male speech, man speaking": 66, "Crowd": 67, "Screech": 68, "Animal": 69, "Human group actions": 70, "Telephone": 71, "Tools": 72, "Giggle": 73, "Crushing": 74, "Thump, thud": 75, "Hammer": 76, "Engine": 77, "Cupboard open or close": 78, "Glass": 79, "Writing": 80, "Clock": 81, "Plucked string instrument": 82, "Fowl": 83, "Water tap, faucet": 84, "Knock": 85, "Trickle, dribble": 86, "Rattle": 87, "Conversation": 88, "Accelerating, revving, vroom": 89, "Fixed-wing aircraft, airplane": 90, "Screaming": 91, "Walk, footsteps": 92, "Stream": 93, "Printer": 94, "Traffic noise, roadway noise": 95, "Motorcycle": 96, "Water": 97, "Scratching (performance technique)": 98, "Tap": 99, "Percussion": 100, "Chuckle, chortle": 101, "Motor vehicle (road)": 102, "Crow": 103, "Vehicle horn, car horn, honking": 104, "Bird vocalization, bird call, bird song": 105, "Drill": 106, "Race car, auto racing": 107, "Meow": 108, "Bass drum": 109, "Drum kit": 110, "Wild animals": 111, "Crash cymbal": 112, "Cough": 113, "Typing": 114, "Bowed string instrument": 115, "Computer keyboard": 116, "Vehicle": 117, "Train": 118, "Applause": 119, "Bicycle": 120, "Tick": 121, "Drum": 122, "Burping, eructation": 123, "Bicycle bell": 124, "Cowbell": 125, "Accordion": 126, "Toilet flush": 127, "Purr": 128, "Church bell": 129, "Cat": 130, "Insect": 131, "Engine starting": 132, "Chewing, mastication": 133, "Sink (filling or washing)": 134, "Dog": 135, "Bird": 136, "Finger snapping": 137, "Child speech, kid speaking": 138, "Wood": 139, "Music": 140, "Sawing": 141, "Bell": 142, "Fireworks": 143, "Crumpling, crinkling": 144, "Ocean": 145, "Gurgling": 146, "Fart": 147, "Mechanisms": 148, "Acoustic guitar": 149, "Singing": 150, "Boom": 151, "Bus": 152, "Cutlery, silverware": 153, "Liquid": 154, "Explosion": 155, "Gull, seagull": 156, "Thunder": 157, "Siren": 158, "Marimba, xylophone": 159, "Female singing": 160, "Tick-tock": 161, "Frog": 162, "Frying (food)": 163, "Buzz": 164, "Car passing by": 165, "Electric guitar": 166, "Gasp": 167, "Rattle (instrument)": 168, "Piano": 169, "Doorbell": 170, "Chime": 171, "Car": 172, "Fire": 173, "Trumpet": 174, "Truck": 175, "Hands": 176, "Domestic animals, pets": 177, "Chirp, tweet": 178, "Breathing": 179, "Cricket": 180, "Tambourine": 181, "Bass guitar": 182, "Idling": 183, "Scissors": 184, "Rain": 185, "Strum": 186, "Shout": 187, "Keys jangling": 188, "Camera": 189, "Hiss": 190, "Growling": 191, "Snare drum": 192, "Brass instrument": 193, "Bathtub (filling or washing)": 194, "Typewriter": 195, "Aircraft": 196, "Organ": 197, "Laughter": 198, "Harmonica": 199}
\ No newline at end of file
diff --git a/class_labels/UrbanSound8K_class_labels_indices.json b/class_labels/UrbanSound8K_class_labels_indices.json
new file mode 100644
index 0000000..496f37f
--- /dev/null
+++ b/class_labels/UrbanSound8K_class_labels_indices.json
@@ -0,0 +1,1 @@
+{"air conditioner": 0, "car horn": 1, "children playing": 2, "dog bark": 3, "drilling": 4, "engine idling": 5, "gun shot": 6, "jackhammer": 7, "siren": 8, "street music": 9}
\ No newline at end of file
diff --git a/class_labels/VGGSound_class_labels_indices.json b/class_labels/VGGSound_class_labels_indices.json
new file mode 100644
index 0000000..d794b2c
--- /dev/null
+++ b/class_labels/VGGSound_class_labels_indices.json
@@ -0,0 +1,1 @@
+{"people crowd": 0, "playing mandolin": 1, "pumping water": 2, "horse neighing": 3, "airplane flyby": 4, "playing drum kit": 5, "pheasant crowing": 6, "duck quacking": 7, "wood thrush calling": 8, "dog bow-wow": 9, "arc welding": 10, "writing on blackboard with chalk": 11, "forging swords": 12, "swimming": 13, "bee, wasp, etc. buzzing": 14, "child singing": 15, "mouse clicking": 16, "playing trombone": 17, "telephone bell ringing": 18, "beat boxing": 19, "cattle mooing": 20, "lions roaring": 21, "ambulance siren": 22, "gibbon howling": 23, "people sniggering": 24, "playing clarinet": 25, "playing bassoon": 26, "playing bongo": 27, "playing electric guitar": 28, "playing badminton": 29, "bull bellowing": 30, "cat caterwauling": 31, "playing sitar": 32, "whale calling": 33, "snake hissing": 34, "people burping": 35, "francolin calling": 36, "fireworks banging": 37, "driving buses": 38, "people belly laughing": 39, "chicken clucking": 40, "playing double bass": 41, "canary calling": 42, "people battle cry": 43, "male singing": 44, "horse clip-clop": 45, "baby crying": 46, "cow lowing": 47, "reversing beeps": 48, "otter growling": 49, "cheetah chirrup": 50, "people running": 51, "ice cream truck, ice cream van": 52, "playing harpsichord": 53, "heart sounds, heartbeat": 54, "pig oinking": 55, "police radio chatter": 56, "cat hissing": 57, "wind chime": 58, "elk bugling": 59, "lions growling": 60, "fly, housefly buzzing": 61, "ferret dooking": 62, "railroad car, train wagon": 63, "church bell ringing": 64, "cat meowing": 65, "wind rustling leaves": 66, "bouncing on trampoline": 67, "mouse squeaking": 68, "sheep bleating": 69, "people eating crisps": 70, "people sneezing": 71, "playing squash": 72, "footsteps on snow": 73, "people humming": 74, "tap dancing": 75, "snake rattling": 76, "elephant trumpeting": 77, "people booing": 78, "disc scratching": 79, "skidding": 80, "cupboard opening or closing": 81, "playing bagpipes": 82, "basketball bounce": 83, "chinchilla barking": 84, "parrot talking": 85, "woodpecker pecking tree": 86, "fire truck siren": 87, "slot machine": 88, "playing french horn": 89, "air conditioning noise": 90, "people finger snapping": 91, "eagle screaming": 92, "playing harmonica": 93, "playing tympani": 94, "zebra braying": 95, "hedge trimmer running": 96, "playing acoustic guitar": 97, "hair dryer drying": 98, "orchestra": 99, "playing darts": 100, "children shouting": 101, "people slurping": 102, "alligators, crocodiles hissing": 103, "mouse pattering": 104, "people marching": 105, "vehicle horn, car horn, honking": 106, "sea lion barking": 107, "people clapping": 108, "hail": 109, "fire crackling": 110, "bathroom ventilation fan running": 111, "opening or closing car doors": 112, "skiing": 113, "dog barking": 114, "race car, auto racing": 115, "subway, metro, underground": 116, "underwater bubbling": 117, "car passing by": 118, "playing tennis": 119, "warbler chirping": 120, "helicopter": 121, "driving motorcycle": 122, "train wheels squealing": 123, "baby laughter": 124, "driving snowmobile": 125, "bird squawking": 126, "cuckoo bird calling": 127, "people whistling": 128, "shot football": 129, "playing tuning fork": 130, "dog howling": 131, "playing violin, fiddle": 132, "people eating": 133, "baltimore oriole calling": 134, "playing timbales": 135, "door slamming": 136, "people shuffling": 137, "typing on typewriter": 138, "magpie calling": 139, "playing harp": 140, "playing hammond organ": 141, "people eating apple": 142, "mosquito buzzing": 143, "playing oboe": 144, "playing volleyball": 145, "using sewing machines": 146, "electric grinder grinding": 147, "cutting hair with electric trimmers": 148, "splashing water": 149, "people sobbing": 150, "female singing": 151, "wind noise": 152, "car engine knocking": 153, "black capped chickadee calling": 154, "people screaming": 155, "cat growling": 156, "penguins braying": 157, "people coughing": 158, "metronome": 159, "train horning": 160, "goat bleating": 161, "playing tambourine": 162, "fox barking": 163, "airplane": 164, "firing cannon": 165, "thunder": 166, "smoke detector beeping": 167, "playing erhu": 168, "ice cracking": 169, "dog growling": 170, "playing saxophone": 171, "owl hooting": 172, "playing trumpet": 173, "sailing": 174, "waterfall burbling": 175, "machine gun shooting": 176, "baby babbling": 177, "playing synthesizer": 178, "donkey, ass braying": 179, "people cheering": 180, "playing shofar": 181, "playing hockey": 182, "playing banjo": 183, "cricket chirping": 184, "playing snare drum": 185, "ripping paper": 186, "child speech, kid speaking": 187, "crow cawing": 188, "sloshing water": 189, "playing zither": 190, "scuba diving": 191, "playing steelpan": 192, "goose honking": 193, "tapping guitar": 194, "spraying water": 195, "playing bass drum": 196, "printer printing": 197, "playing ukulele": 198, "ocean burbling": 199, "playing didgeridoo": 200, "sharpen knife": 201, "typing on computer keyboard": 202, "playing table tennis": 203, "rope skipping": 204, "playing marimba, xylophone": 205, "playing bugle": 206, "playing guiro": 207, "playing flute": 208, "tornado roaring": 209, "stream burbling": 210, "electric shaver, electric razor shaving": 211, "playing gong": 212, "eating with cutlery": 213, "playing piano": 214, "people giggling": 215, "chicken crowing": 216, "female speech, woman speaking": 217, "golf driving": 218, "frog croaking": 219, "people eating noodle": 220, "mynah bird singing": 221, "playing timpani": 222, "playing congas": 223, "dinosaurs bellowing": 224, "playing bass guitar": 225, "turkey gobbling": 226, "chipmunk chirping": 227, "chopping food": 228, "striking bowling": 229, "missile launch": 230, "squishing water": 231, "civil defense siren": 232, "blowtorch igniting": 233, "tractor digging": 234, "lighting firecrackers": 235, "playing theremin": 236, "train whistling": 237, "people nose blowing": 238, "car engine starting": 239, "lathe spinning": 240, "playing cello": 241, "motorboat, speedboat acceleration": 242, "playing vibraphone": 243, "playing washboard": 244, "playing cornet": 245, "pigeon, dove cooing": 246, "roller coaster running": 247, "opening or closing car electric windows": 248, "foghorn": 249, "coyote howling": 250, "hammering nails": 251, "toilet flushing": 252, "strike lighter": 253, "bird wings flapping": 254, "playing steel guitar, slide guitar": 255, "volcano explosion": 256, "people whispering": 257, "bowling impact": 258, "yodelling": 259, "firing muskets": 260, "raining": 261, "singing bowl": 262, "plastic bottle crushing": 263, "chimpanzee pant-hooting": 264, "playing electronic organ": 265, "chainsawing trees": 266, "dog baying": 267, "lawn mowing": 268, "people babbling": 269, "striking pool": 270, "eletric blender running": 271, "playing tabla": 272, "cap gun shooting": 273, "planing timber": 274, "air horn": 275, "sliding door": 276, "cell phone buzzing": 277, "sea waves": 278, "playing castanets": 279, "singing choir": 280, "people slapping": 281, "barn swallow calling": 282, "people hiccup": 283, "vacuum cleaner cleaning floors": 284, "playing lacrosse": 285, "bird chirping, tweeting": 286, "lip smacking": 287, "chopping wood": 288, "police car (siren)": 289, "running electric fan": 290, "cattle, bovinae cowbell": 291, "people gargling": 292, "opening or closing drawers": 293, "playing djembe": 294, "skateboarding": 295, "cat purring": 296, "rowboat, canoe, kayak rowing": 297, "engine accelerating, revving, vroom": 298, "playing glockenspiel": 299, "popping popcorn": 300, "car engine idling": 301, "alarm clock ringing": 302, "dog whimpering": 303, "playing accordion": 304, "playing cymbal": 305, "male speech, man speaking": 306, "rapping": 307, "people farting": 308}
\ No newline at end of file
diff --git a/class_labels/audioset_class_labels_indices.json b/class_labels/audioset_class_labels_indices.json
new file mode 100644
index 0000000..0f6780b
--- /dev/null
+++ b/class_labels/audioset_class_labels_indices.json
@@ -0,0 +1,1 @@
+{"Speech": 0, "Male speech, man speaking": 1, "Female speech, woman speaking": 2, "Child speech, kid speaking": 3, "Conversation": 4, "Narration, monologue": 5, "Babbling": 6, "Speech synthesizer": 7, "Shout": 8, "Bellow": 9, "Whoop": 10, "Yell": 11, "Battle cry": 12, "Children shouting": 13, "Screaming": 14, "Whispering": 15, "Laughter": 16, "Baby laughter": 17, "Giggle": 18, "Snicker": 19, "Belly laugh": 20, "Chuckle, chortle": 21, "Crying, sobbing": 22, "Baby cry, infant cry": 23, "Whimper": 24, "Wail, moan": 25, "Sigh": 26, "Singing": 27, "Choir": 28, "Yodeling": 29, "Chant": 30, "Mantra": 31, "Male singing": 32, "Female singing": 33, "Child singing": 34, "Synthetic singing": 35, "Rapping": 36, "Humming": 37, "Groan": 38, "Grunt": 39, "Whistling": 40, "Breathing": 41, "Wheeze": 42, "Snoring": 43, "Gasp": 44, "Pant": 45, "Snort": 46, "Cough": 47, "Throat clearing": 48, "Sneeze": 49, "Sniff": 50, "Run": 51, "Shuffle": 52, "Walk, footsteps": 53, "Chewing, mastication": 54, "Biting": 55, "Gargling": 56, "Stomach rumble": 57, "Burping, eructation": 58, "Hiccup": 59, "Fart": 60, "Hands": 61, "Finger snapping": 62, "Clapping": 63, "Heart sounds, heartbeat": 64, "Heart murmur": 65, "Cheering": 66, "Applause": 67, "Chatter": 68, "Crowd": 69, "Hubbub, speech noise, speech babble": 70, "Children playing": 71, "Animal": 72, "Domestic animals, pets": 73, "Dog": 74, "Bark": 75, "Yip": 76, "Howl": 77, "Bow-wow": 78, "Growling": 79, "Whimper (dog)": 80, "Cat": 81, "Purr": 82, "Meow": 83, "Hiss": 84, "Caterwaul": 85, "Livestock, farm animals, working animals": 86, "Horse": 87, "Clip-clop": 88, "Neigh, whinny": 89, "Cattle, bovinae": 90, "Moo": 91, "Cowbell": 92, "Pig": 93, "Oink": 94, "Goat": 95, "Bleat": 96, "Sheep": 97, "Fowl": 98, "Chicken, rooster": 99, "Cluck": 100, "Crowing, cock-a-doodle-doo": 101, "Turkey": 102, "Gobble": 103, "Duck": 104, "Quack": 105, "Goose": 106, "Honk": 107, "Wild animals": 108, "Roaring cats (lions, tigers)": 109, "Roar": 110, "Bird": 111, "Bird vocalization, bird call, bird song": 112, "Chirp, tweet": 113, "Squawk": 114, "Pigeon, dove": 115, "Coo": 116, "Crow": 117, "Caw": 118, "Owl": 119, "Hoot": 120, "Bird flight, flapping wings": 121, "Canidae, dogs, wolves": 122, "Rodents, rats, mice": 123, "Mouse": 124, "Patter": 125, "Insect": 126, "Cricket": 127, "Mosquito": 128, "Fly, housefly": 129, "Buzz": 130, "Bee, wasp, etc.": 131, "Frog": 132, "Croak": 133, "Snake": 134, "Rattle": 135, "Whale vocalization": 136, "Music": 137, "Musical instrument": 138, "Plucked string instrument": 139, "Guitar": 140, "Electric guitar": 141, "Bass guitar": 142, "Acoustic guitar": 143, "Steel guitar, slide guitar": 144, "Tapping (guitar technique)": 145, "Strum": 146, "Banjo": 147, "Sitar": 148, "Mandolin": 149, "Zither": 150, "Ukulele": 151, "Keyboard (musical)": 152, "Piano": 153, "Electric piano": 154, "Organ": 155, "Electronic organ": 156, "Hammond organ": 157, "Synthesizer": 158, "Sampler": 159, "Harpsichord": 160, "Percussion": 161, "Drum kit": 162, "Drum machine": 163, "Drum": 164, "Snare drum": 165, "Rimshot": 166, "Drum roll": 167, "Bass drum": 168, "Timpani": 169, "Tabla": 170, "Cymbal": 171, "Hi-hat": 172, "Wood block": 173, "Tambourine": 174, "Rattle (instrument)": 175, "Maraca": 176, "Gong": 177, "Tubular bells": 178, "Mallet percussion": 179, "Marimba, xylophone": 180, "Glockenspiel": 181, "Vibraphone": 182, "Steelpan": 183, "Orchestra": 184, "Brass instrument": 185, "French horn": 186, "Trumpet": 187, "Trombone": 188, "Bowed string instrument": 189, "String section": 190, "Violin, fiddle": 191, "Pizzicato": 192, "Cello": 193, "Double bass": 194, "Wind instrument, woodwind instrument": 195, "Flute": 196, "Saxophone": 197, "Clarinet": 198, "Harp": 199, "Bell": 200, "Church bell": 201, "Jingle bell": 202, "Bicycle bell": 203, "Tuning fork": 204, "Chime": 205, "Wind chime": 206, "Change ringing (campanology)": 207, "Harmonica": 208, "Accordion": 209, "Bagpipes": 210, "Didgeridoo": 211, "Shofar": 212, "Theremin": 213, "Singing bowl": 214, "Scratching (performance technique)": 215, "Pop music": 216, "Hip hop music": 217, "Beatboxing": 218, "Rock music": 219, "Heavy metal": 220, "Punk rock": 221, "Grunge": 222, "Progressive rock": 223, "Rock and roll": 224, "Psychedelic rock": 225, "Rhythm and blues": 226, "Soul music": 227, "Reggae": 228, "Country": 229, "Swing music": 230, "Bluegrass": 231, "Funk": 232, "Folk music": 233, "Middle Eastern music": 234, "Jazz": 235, "Disco": 236, "Classical music": 237, "Opera": 238, "Electronic music": 239, "House music": 240, "Techno": 241, "Dubstep": 242, "Drum and bass": 243, "Electronica": 244, "Electronic dance music": 245, "Ambient music": 246, "Trance music": 247, "Music of Latin America": 248, "Salsa music": 249, "Flamenco": 250, "Blues": 251, "Music for children": 252, "New-age music": 253, "Vocal music": 254, "A capella": 255, "Music of Africa": 256, "Afrobeat": 257, "Christian music": 258, "Gospel music": 259, "Music of Asia": 260, "Carnatic music": 261, "Music of Bollywood": 262, "Ska": 263, "Traditional music": 264, "Independent music": 265, "Song": 266, "Background music": 267, "Theme music": 268, "Jingle (music)": 269, "Soundtrack music": 270, "Lullaby": 271, "Video game music": 272, "Christmas music": 273, "Dance music": 274, "Wedding music": 275, "Happy music": 276, "Funny music": 277, "Sad music": 278, "Tender music": 279, "Exciting music": 280, "Angry music": 281, "Scary music": 282, "Wind": 283, "Rustling leaves": 284, "Wind noise (microphone)": 285, "Thunderstorm": 286, "Thunder": 287, "Water": 288, "Rain": 289, "Raindrop": 290, "Rain on surface": 291, "Stream": 292, "Waterfall": 293, "Ocean": 294, "Waves, surf": 295, "Steam": 296, "Gurgling": 297, "Fire": 298, "Crackle": 299, "Vehicle": 300, "Boat, Water vehicle": 301, "Sailboat, sailing ship": 302, "Rowboat, canoe, kayak": 303, "Motorboat, speedboat": 304, "Ship": 305, "Motor vehicle (road)": 306, "Car": 307, "Vehicle horn, car horn, honking": 308, "Toot": 309, "Car alarm": 310, "Power windows, electric windows": 311, "Skidding": 312, "Tire squeal": 313, "Car passing by": 314, "Race car, auto racing": 315, "Truck": 316, "Air brake": 317, "Air horn, truck horn": 318, "Reversing beeps": 319, "Ice cream truck, ice cream van": 320, "Bus": 321, "Emergency vehicle": 322, "Police car (siren)": 323, "Ambulance (siren)": 324, "Fire engine, fire truck (siren)": 325, "Motorcycle": 326, "Traffic noise, roadway noise": 327, "Rail transport": 328, "Train": 329, "Train whistle": 330, "Train horn": 331, "Railroad car, train wagon": 332, "Train wheels squealing": 333, "Subway, metro, underground": 334, "Aircraft": 335, "Aircraft engine": 336, "Jet engine": 337, "Propeller, airscrew": 338, "Helicopter": 339, "Fixed-wing aircraft, airplane": 340, "Bicycle": 341, "Skateboard": 342, "Engine": 343, "Light engine (high frequency)": 344, "Dental drill, dentist's drill": 345, "Lawn mower": 346, "Chainsaw": 347, "Medium engine (mid frequency)": 348, "Heavy engine (low frequency)": 349, "Engine knocking": 350, "Engine starting": 351, "Idling": 352, "Accelerating, revving, vroom": 353, "Door": 354, "Doorbell": 355, "Ding-dong": 356, "Sliding door": 357, "Slam": 358, "Knock": 359, "Tap": 360, "Squeak": 361, "Cupboard open or close": 362, "Drawer open or close": 363, "Dishes, pots, and pans": 364, "Cutlery, silverware": 365, "Chopping (food)": 366, "Frying (food)": 367, "Microwave oven": 368, "Blender": 369, "Water tap, faucet": 370, "Sink (filling or washing)": 371, "Bathtub (filling or washing)": 372, "Hair dryer": 373, "Toilet flush": 374, "Toothbrush": 375, "Electric toothbrush": 376, "Vacuum cleaner": 377, "Zipper (clothing)": 378, "Keys jangling": 379, "Coin (dropping)": 380, "Scissors": 381, "Electric shaver, electric razor": 382, "Shuffling cards": 383, "Typing": 384, "Typewriter": 385, "Computer keyboard": 386, "Writing": 387, "Alarm": 388, "Telephone": 389, "Telephone bell ringing": 390, "Ringtone": 391, "Telephone dialing, DTMF": 392, "Dial tone": 393, "Busy signal": 394, "Alarm clock": 395, "Siren": 396, "Civil defense siren": 397, "Buzzer": 398, "Smoke detector, smoke alarm": 399, "Fire alarm": 400, "Foghorn": 401, "Whistle": 402, "Steam whistle": 403, "Mechanisms": 404, "Ratchet, pawl": 405, "Clock": 406, "Tick": 407, "Tick-tock": 408, "Gears": 409, "Pulleys": 410, "Sewing machine": 411, "Mechanical fan": 412, "Air conditioning": 413, "Cash register": 414, "Printer": 415, "Camera": 416, "Single-lens reflex camera": 417, "Tools": 418, "Hammer": 419, "Jackhammer": 420, "Sawing": 421, "Filing (rasp)": 422, "Sanding": 423, "Power tool": 424, "Drill": 425, "Explosion": 426, "Gunshot, gunfire": 427, "Machine gun": 428, "Fusillade": 429, "Artillery fire": 430, "Cap gun": 431, "Fireworks": 432, "Firecracker": 433, "Burst, pop": 434, "Eruption": 435, "Boom": 436, "Wood": 437, "Chop": 438, "Splinter": 439, "Crack": 440, "Glass": 441, "Chink, clink": 442, "Shatter": 443, "Liquid": 444, "Splash, splatter": 445, "Slosh": 446, "Squish": 447, "Drip": 448, "Pour": 449, "Trickle, dribble": 450, "Gush": 451, "Fill (with liquid)": 452, "Spray": 453, "Pump (liquid)": 454, "Stir": 455, "Boiling": 456, "Sonar": 457, "Arrow": 458, "Whoosh, swoosh, swish": 459, "Thump, thud": 460, "Thunk": 461, "Electronic tuner": 462, "Effects unit": 463, "Chorus effect": 464, "Basketball bounce": 465, "Bang": 466, "Slap, smack": 467, "Whack, thwack": 468, "Smash, crash": 469, "Breaking": 470, "Bouncing": 471, "Whip": 472, "Flap": 473, "Scratch": 474, "Scrape": 475, "Rub": 476, "Roll": 477, "Crushing": 478, "Crumpling, crinkling": 479, "Tearing": 480, "Beep, bleep": 481, "Ping": 482, "Ding": 483, "Clang": 484, "Squeal": 485, "Creak": 486, "Rustle": 487, "Whir": 488, "Clatter": 489, "Sizzle": 490, "Clicking": 491, "Clickety-clack": 492, "Rumble": 493, "Plop": 494, "Jingle, tinkle": 495, "Hum": 496, "Zing": 497, "Boing": 498, "Crunch": 499, "Silence": 500, "Sine wave": 501, "Harmonic": 502, "Chirp tone": 503, "Sound effect": 504, "Pulse": 505, "Inside, small room": 506, "Inside, large room or hall": 507, "Inside, public space": 508, "Outside, urban or manmade": 509, "Outside, rural or natural": 510, "Reverberation": 511, "Echo": 512, "Noise": 513, "Environmental noise": 514, "Static": 515, "Mains hum": 516, "Distortion": 517, "Sidetone": 518, "Cacophony": 519, "White noise": 520, "Pink noise": 521, "Throbbing": 522, "Vibration": 523, "Television": 524, "Radio": 525, "Field recording": 526}
\ No newline at end of file
diff --git a/class_labels/audioset_fsd50k_class_labels_indices.json b/class_labels/audioset_fsd50k_class_labels_indices.json
new file mode 100644
index 0000000..fd8014e
--- /dev/null
+++ b/class_labels/audioset_fsd50k_class_labels_indices.json
@@ -0,0 +1,1 @@
+{"Speech": 0, "Male speech, man speaking": 1, "Female speech, woman speaking": 2, "Child speech, kid speaking": 3, "Conversation": 4, "Narration, monologue": 5, "Babbling": 6, "Speech synthesizer": 7, "Shout": 8, "Bellow": 9, "Whoop": 10, "Yell": 11, "Battle cry": 12, "Children shouting": 13, "Screaming": 14, "Whispering": 15, "Laughter": 16, "Baby laughter": 17, "Giggle": 18, "Snicker": 19, "Belly laugh": 20, "Chuckle, chortle": 21, "Crying, sobbing": 22, "Baby cry, infant cry": 23, "Whimper": 24, "Wail, moan": 25, "Sigh": 26, "Singing": 27, "Choir": 28, "Yodeling": 29, "Chant": 30, "Mantra": 31, "Male singing": 32, "Female singing": 33, "Child singing": 34, "Synthetic singing": 35, "Rapping": 36, "Humming": 37, "Groan": 38, "Grunt": 39, "Whistling": 40, "Breathing": 41, "Wheeze": 42, "Snoring": 43, "Gasp": 44, "Pant": 45, "Snort": 46, "Cough": 47, "Throat clearing": 48, "Sneeze": 49, "Sniff": 50, "Run": 51, "Shuffle": 52, "Walk, footsteps": 53, "Chewing, mastication": 54, "Biting": 55, "Gargling": 56, "Stomach rumble": 57, "Burping, eructation": 58, "Hiccup": 59, "Fart": 60, "Hands": 61, "Finger snapping": 62, "Clapping": 63, "Heart sounds, heartbeat": 64, "Heart murmur": 65, "Cheering": 66, "Applause": 67, "Chatter": 68, "Crowd": 69, "Hubbub, speech noise, speech babble": 70, "Children playing": 71, "Animal": 72, "Domestic animals, pets": 73, "Dog": 74, "Bark": 75, "Yip": 76, "Howl": 77, "Bow-wow": 78, "Growling": 79, "Whimper (dog)": 80, "Cat": 81, "Purr": 82, "Meow": 83, "Hiss": 84, "Caterwaul": 85, "Livestock, farm animals, working animals": 86, "Horse": 87, "Clip-clop": 88, "Neigh, whinny": 89, "Cattle, bovinae": 90, "Moo": 91, "Cowbell": 92, "Pig": 93, "Oink": 94, "Goat": 95, "Bleat": 96, "Sheep": 97, "Fowl": 98, "Chicken, rooster": 99, "Cluck": 100, "Crowing, cock-a-doodle-doo": 101, "Turkey": 102, "Gobble": 103, "Duck": 104, "Quack": 105, "Goose": 106, "Honk": 107, "Wild animals": 108, "Roaring cats (lions, tigers)": 109, "Roar": 110, "Bird": 111, "Bird vocalization, bird call, bird song": 112, "Chirp, tweet": 113, "Squawk": 114, "Pigeon, dove": 115, "Coo": 116, "Crow": 117, "Caw": 118, "Owl": 119, "Hoot": 120, "Bird flight, flapping wings": 121, "Canidae, dogs, wolves": 122, "Rodents, rats, mice": 123, "Mouse": 124, "Patter": 125, "Insect": 126, "Cricket": 127, "Mosquito": 128, "Fly, housefly": 129, "Buzz": 130, "Bee, wasp, etc.": 131, "Frog": 132, "Croak": 133, "Snake": 134, "Rattle": 135, "Whale vocalization": 136, "Music": 137, "Musical instrument": 138, "Plucked string instrument": 139, "Guitar": 140, "Electric guitar": 141, "Bass guitar": 142, "Acoustic guitar": 143, "Steel guitar, slide guitar": 144, "Tapping (guitar technique)": 145, "Strum": 146, "Banjo": 147, "Sitar": 148, "Mandolin": 149, "Zither": 150, "Ukulele": 151, "Keyboard (musical)": 152, "Piano": 153, "Electric piano": 154, "Organ": 155, "Electronic organ": 156, "Hammond organ": 157, "Synthesizer": 158, "Sampler": 159, "Harpsichord": 160, "Percussion": 161, "Drum kit": 162, "Drum machine": 163, "Drum": 164, "Snare drum": 165, "Rimshot": 166, "Drum roll": 167, "Bass drum": 168, "Timpani": 169, "Tabla": 170, "Cymbal": 171, "Hi-hat": 172, "Wood block": 173, "Tambourine": 174, "Rattle (instrument)": 175, "Maraca": 176, "Gong": 177, "Tubular bells": 178, "Mallet percussion": 179, "Marimba, xylophone": 180, "Glockenspiel": 181, "Vibraphone": 182, "Steelpan": 183, "Orchestra": 184, "Brass instrument": 185, "French horn": 186, "Trumpet": 187, "Trombone": 188, "Bowed string instrument": 189, "String section": 190, "Violin, fiddle": 191, "Pizzicato": 192, "Cello": 193, "Double bass": 194, "Wind instrument, woodwind instrument": 195, "Flute": 196, "Saxophone": 197, "Clarinet": 198, "Harp": 199, "Bell": 200, "Church bell": 201, "Jingle bell": 202, "Bicycle bell": 203, "Tuning fork": 204, "Chime": 205, "Wind chime": 206, "Change ringing (campanology)": 207, "Harmonica": 208, "Accordion": 209, "Bagpipes": 210, "Didgeridoo": 211, "Shofar": 212, "Theremin": 213, "Singing bowl": 214, "Scratching (performance technique)": 215, "Pop music": 216, "Hip hop music": 217, "Beatboxing": 218, "Rock music": 219, "Heavy metal": 220, "Punk rock": 221, "Grunge": 222, "Progressive rock": 223, "Rock and roll": 224, "Psychedelic rock": 225, "Rhythm and blues": 226, "Soul music": 227, "Reggae": 228, "Country": 229, "Swing music": 230, "Bluegrass": 231, "Funk": 232, "Folk music": 233, "Middle Eastern music": 234, "Jazz": 235, "Disco": 236, "Classical music": 237, "Opera": 238, "Electronic music": 239, "House music": 240, "Techno": 241, "Dubstep": 242, "Drum and bass": 243, "Electronica": 244, "Electronic dance music": 245, "Ambient music": 246, "Trance music": 247, "Music of Latin America": 248, "Salsa music": 249, "Flamenco": 250, "Blues": 251, "Music for children": 252, "New-age music": 253, "Vocal music": 254, "A capella": 255, "Music of Africa": 256, "Afrobeat": 257, "Christian music": 258, "Gospel music": 259, "Music of Asia": 260, "Carnatic music": 261, "Music of Bollywood": 262, "Ska": 263, "Traditional music": 264, "Independent music": 265, "Song": 266, "Background music": 267, "Theme music": 268, "Jingle (music)": 269, "Soundtrack music": 270, "Lullaby": 271, "Video game music": 272, "Christmas music": 273, "Dance music": 274, "Wedding music": 275, "Happy music": 276, "Funny music": 277, "Sad music": 278, "Tender music": 279, "Exciting music": 280, "Angry music": 281, "Scary music": 282, "Wind": 283, "Rustling leaves": 284, "Wind noise (microphone)": 285, "Thunderstorm": 286, "Thunder": 287, "Water": 288, "Rain": 289, "Raindrop": 290, "Rain on surface": 291, "Stream": 292, "Waterfall": 293, "Ocean": 294, "Waves, surf": 295, "Steam": 296, "Gurgling": 297, "Fire": 298, "Crackle": 299, "Vehicle": 300, "Boat, Water vehicle": 301, "Sailboat, sailing ship": 302, "Rowboat, canoe, kayak": 303, "Motorboat, speedboat": 304, "Ship": 305, "Motor vehicle (road)": 306, "Car": 307, "Vehicle horn, car horn, honking": 308, "Toot": 309, "Car alarm": 310, "Power windows, electric windows": 311, "Skidding": 312, "Tire squeal": 313, "Car passing by": 314, "Race car, auto racing": 315, "Truck": 316, "Air brake": 317, "Air horn, truck horn": 318, "Reversing beeps": 319, "Ice cream truck, ice cream van": 320, "Bus": 321, "Emergency vehicle": 322, "Police car (siren)": 323, "Ambulance (siren)": 324, "Fire engine, fire truck (siren)": 325, "Motorcycle": 326, "Traffic noise, roadway noise": 327, "Rail transport": 328, "Train": 329, "Train whistle": 330, "Train horn": 331, "Railroad car, train wagon": 332, "Train wheels squealing": 333, "Subway, metro, underground": 334, "Aircraft": 335, "Aircraft engine": 336, "Jet engine": 337, "Propeller, airscrew": 338, "Helicopter": 339, "Fixed-wing aircraft, airplane": 340, "Bicycle": 341, "Skateboard": 342, "Engine": 343, "Light engine (high frequency)": 344, "Dental drill, dentist's drill": 345, "Lawn mower": 346, "Chainsaw": 347, "Medium engine (mid frequency)": 348, "Heavy engine (low frequency)": 349, "Engine knocking": 350, "Engine starting": 351, "Idling": 352, "Accelerating, revving, vroom": 353, "Door": 354, "Doorbell": 355, "Ding-dong": 356, "Sliding door": 357, "Slam": 358, "Knock": 359, "Tap": 360, "Squeak": 361, "Cupboard open or close": 362, "Drawer open or close": 363, "Dishes, pots, and pans": 364, "Cutlery, silverware": 365, "Chopping (food)": 366, "Frying (food)": 367, "Microwave oven": 368, "Blender": 369, "Water tap, faucet": 370, "Sink (filling or washing)": 371, "Bathtub (filling or washing)": 372, "Hair dryer": 373, "Toilet flush": 374, "Toothbrush": 375, "Electric toothbrush": 376, "Vacuum cleaner": 377, "Zipper (clothing)": 378, "Keys jangling": 379, "Coin (dropping)": 380, "Scissors": 381, "Electric shaver, electric razor": 382, "Shuffling cards": 383, "Typing": 384, "Typewriter": 385, "Computer keyboard": 386, "Writing": 387, "Alarm": 388, "Telephone": 389, "Telephone bell ringing": 390, "Ringtone": 391, "Telephone dialing, DTMF": 392, "Dial tone": 393, "Busy signal": 394, "Alarm clock": 395, "Siren": 396, "Civil defense siren": 397, "Buzzer": 398, "Smoke detector, smoke alarm": 399, "Fire alarm": 400, "Foghorn": 401, "Whistle": 402, "Steam whistle": 403, "Mechanisms": 404, "Ratchet, pawl": 405, "Clock": 406, "Tick": 407, "Tick-tock": 408, "Gears": 409, "Pulleys": 410, "Sewing machine": 411, "Mechanical fan": 412, "Air conditioning": 413, "Cash register": 414, "Printer": 415, "Camera": 416, "Single-lens reflex camera": 417, "Tools": 418, "Hammer": 419, "Jackhammer": 420, "Sawing": 421, "Filing (rasp)": 422, "Sanding": 423, "Power tool": 424, "Drill": 425, "Explosion": 426, "Gunshot, gunfire": 427, "Machine gun": 428, "Fusillade": 429, "Artillery fire": 430, "Cap gun": 431, "Fireworks": 432, "Firecracker": 433, "Burst, pop": 434, "Eruption": 435, "Boom": 436, "Wood": 437, "Chop": 438, "Splinter": 439, "Crack": 440, "Glass": 441, "Chink, clink": 442, "Shatter": 443, "Liquid": 444, "Splash, splatter": 445, "Slosh": 446, "Squish": 447, "Drip": 448, "Pour": 449, "Trickle, dribble": 450, "Gush": 451, "Fill (with liquid)": 452, "Spray": 453, "Pump (liquid)": 454, "Stir": 455, "Boiling": 456, "Sonar": 457, "Arrow": 458, "Whoosh, swoosh, swish": 459, "Thump, thud": 460, "Thunk": 461, "Electronic tuner": 462, "Effects unit": 463, "Chorus effect": 464, "Basketball bounce": 465, "Bang": 466, "Slap, smack": 467, "Whack, thwack": 468, "Smash, crash": 469, "Breaking": 470, "Bouncing": 471, "Whip": 472, "Flap": 473, "Scratch": 474, "Scrape": 475, "Rub": 476, "Roll": 477, "Crushing": 478, "Crumpling, crinkling": 479, "Tearing": 480, "Beep, bleep": 481, "Ping": 482, "Ding": 483, "Clang": 484, "Squeal": 485, "Creak": 486, "Rustle": 487, "Whir": 488, "Clatter": 489, "Sizzle": 490, "Clicking": 491, "Clickety-clack": 492, "Rumble": 493, "Plop": 494, "Jingle, tinkle": 495, "Hum": 496, "Zing": 497, "Boing": 498, "Crunch": 499, "Silence": 500, "Sine wave": 501, "Harmonic": 502, "Chirp tone": 503, "Sound effect": 504, "Pulse": 505, "Inside, small room": 506, "Inside, large room or hall": 507, "Inside, public space": 508, "Outside, urban or manmade": 509, "Outside, rural or natural": 510, "Reverberation": 511, "Echo": 512, "Noise": 513, "Environmental noise": 514, "Static": 515, "Mains hum": 516, "Distortion": 517, "Sidetone": 518, "Cacophony": 519, "White noise": 520, "Pink noise": 521, "Throbbing": 522, "Vibration": 523, "Television": 524, "Radio": 525, "Field recording": 526, "Respiratory sounds": 527, "Human voice": 528, "Packing tape, duct tape": 529, "Domestic sounds, home sounds": 530, "Screech": 531, "Human group actions": 532, "Crash cymbal": 533, "Gull, seagull": 534}
\ No newline at end of file
diff --git a/dev.ipynb b/dev.ipynb
index feeb037..4e0e17d 100644
--- a/dev.ipynb
+++ b/dev.ipynb
@@ -4,14 +4,35 @@
    "cell_type": "code",
    "execution_count": 1,
    "id": "27bf8dca",
-   "metadata": {},
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-08-18T03:53:09.343494Z",
+     "start_time": "2022-08-18T03:53:03.764543Z"
+    }
+   },
    "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "C:\\ProgramData\\Miniconda3\\envs\\py310\\lib\\site-packages\\torchaudio\\compliance\\kaldi.py:22: UserWarning: Failed to initialize NumPy: module compiled against API version 0x10 but this version of numpy is 0xf (Triggered internally at  ..\\torch\\csrc\\utils\\tensor_numpy.cpp:68.)\n",
+      "  EPSILON = torch.tensor(torch.finfo(torch.float).eps)\n"
+     ]
+    },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "1.11.0\n",
-      "0.11.0\n"
+      "1.12.0+cu116\n",
+      "0.12.0+cu116\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "C:\\ProgramData\\Miniconda3\\envs\\py310\\lib\\site-packages\\torchaudio\\backend\\utils.py:62: UserWarning: No audio backend is available.\n",
+      "  warnings.warn(\"No audio backend is available.\")\n"
      ]
     }
    ],
@@ -28,31 +49,13 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "fcafb812",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "id": "aab08868",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import math\n",
-    "import time\n",
-    "\n",
-    "import librosa\n",
-    "import matplotlib.pyplot as plt\n",
-    "import pandas as pd"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
    "id": "e81f3909",
-   "metadata": {},
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-08-18T03:53:10.928313Z",
+     "start_time": "2022-08-18T03:53:10.928279Z"
+    }
+   },
    "outputs": [],
    "source": [
     "DEFAULT_OFFSET = 201\n",
@@ -492,7 +495,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python3.10",
+   "display_name": "python310",
    "language": "python",
    "name": "py310"
   },
@@ -506,7 +509,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.0"
+   "version": "3.10.4"
   }
  },
  "nbformat": 4,
diff --git a/idc_eval.npy b/idc_eval.npy
new file mode 100644
index 0000000..486ef8b
Binary files /dev/null and b/idc_eval.npy differ
diff --git a/idc_train.npy b/idc_train.npy
new file mode 100644
index 0000000..543f01b
Binary files /dev/null and b/idc_train.npy differ
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..573fc11
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,64 @@
+appdirs==1.4.4
+attrs==22.1.0
+audioread==3.0.0
+braceexpand==0.1.7
+certifi==2022.6.15
+cffi==1.15.1
+charset-normalizer==2.1.1
+click==8.1.3
+decorator==5.1.1
+docker-pycreds==0.4.0
+execnet==1.9.0
+ftfy==6.1.1
+gitdb==4.0.9
+GitPython==3.1.27
+h5py==3.7.0
+idna==3.3
+iniconfig==1.1.1
+joblib==1.1.0
+librosa==0.9.2
+llvmlite==0.39.0
+numba==0.56.0
+numpy==1.21.0
+packaging==21.3
+pandas==1.4.4
+pathtools==0.1.2
+Pillow==9.2.0
+pluggy==1.0.0
+pooch==1.6.0
+promise==2.3
+protobuf==3.20.1
+psutil==5.9.1
+py==1.11.0
+pycparser==2.21
+pyparsing==3.0.9
+pytest==7.0.1
+pytest-forked==1.4.0
+pytest-xdist==2.5.0
+python-dateutil==2.8.2
+pytz==2022.2.1
+PyYAML==6.0
+regex==2022.8.17
+requests==2.28.1
+resampy==0.4.0
+scikit-learn==1.1.2
+scipy==1.9.1
+sentry-sdk==1.9.6
+setproctitle==1.3.2
+shortuuid==1.0.9
+six==1.16.0
+smmap==5.0.0
+SoundFile==0.10.3.post1
+threadpoolctl==3.1.0
+tomli==2.0.1
+torch==1.12.1+cu113
+torchaudio==0.12.1+cu113
+torchlibrosa==0.0.9
+torchvision==0.13.1+cu113
+tqdm==4.64.0
+typing-extensions==4.3.0
+urllib3==1.26.12
+wandb==0.13.2
+wcwidth==0.2.5
+webdataset==0.2.20
+wget==3.2
diff --git a/run_code.txt b/run_code.txt
index 70bb6a6..03b9632 100644
--- a/run_code.txt
+++ b/run_code.txt
@@ -14,11 +14,70 @@ CUDA_VISIBLE_DEVICES=0 python -m training.main \
     --epochs=200  \
     --workers=4 \
     --use-bn-sync  \
-    --model PANN-14 \
-    --resample-method="None" \
+    --amodel PANN-14 \
+    --tmodel transformer \
     --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" \
     --datasetinfos "train" "unbalanced_train" "balanced_train"
 
+# Toy dataset Linear Probe
+CUDA_VISIBLE_DEVICES=0 python -m training.main \
+    --save-frequency 1 \
+    --dataset-type="toy" \
+    --train-data="/home/la/kechen/Research/ke_zsasp/workspace/hdf5s/indexes/balanced_train.h5"  \
+    --val-data="/home/la/kechen/Research/ke_zsasp/workspace/hdf5s/indexes/eval.h5"  \
+    --train-ipc="/home/la/kechen/Research/clap/CLAP/idc_train.npy" \
+    --val-ipc="/home/la/kechen/Research/clap/CLAP/idc_eval.npy" \
+    --precision="fp32" \
+    --warmup 10000 \
+    --batch-size=4 \
+    --lr=1e-3 \
+    --wd=0.1 \
+    --epochs=1 \
+    --workers=0 \
+    --amodel HTSAT-tiny \
+    --tmodel transformer \
+    --no-eval \
+    --lp-freeze \
+    --pretrained="/home/la/kechen/Research/clap/ckpt/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_0.pt"
+
+torchrun --nproc_per_node 2 -m training.lp_main \
+    --save-frequency 1 \
+    --dataset-type="toy" \
+    --train-data="/home/la/kechen/Research/ke_zsasp/workspace/hdf5s/indexes/balanced_train.h5"  \
+    --val-data="/home/la/kechen/Research/ke_zsasp/workspace/hdf5s/indexes/eval.h5"  \
+    --train-ipc="/home/la/kechen/Research/clap/CLAP/idc_train.npy" \
+    --val-ipc="/home/la/kechen/Research/clap/CLAP/idc_eval.npy" \
+    --precision="fp32" \
+    --warmup 10000 \
+    --batch-size=32 \
+    --lr=1e-3 \
+    --wd=0.1 \
+    --epochs=1 \
+    --workers=4 \
+    --amodel HTSAT-tiny \
+    --tmodel transformer \
+    --no-eval \
+    --lp-freeze \
+    --pretrained="/home/la/kechen/Research/clap/ckpt/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_0.pt"
+
+CUDA_VISIBLE_DEVICES=0 python -m training.main \
+    --save-frequency 1 \
+    --dataset-type="toy" \
+    --train-data="/home/la/kechen/Research/ke_zsasp/workspace/hdf5s/indexes/balanced_train.h5"  \
+    --val-data="/home/la/kechen/Research/ke_zsasp/workspace/hdf5s/indexes/eval.h5"  \
+    --train-ipc="/home/la/kechen/Research/clap/CLAP/idc_train.npy" \
+    --val-ipc="/home/la/kechen/Research/clap/CLAP/idc_eval.npy" \
+    --precision="fp32" \
+    --warmup 10000 \
+    --batch-size=32 \
+    --lr=1e-3 \
+    --wd=0.1 \
+    --epochs=1 \
+    --workers=4 \
+    --amodel PANN-14 \
+    --tmodel transformer \
+
+
 
 torchrun --nproc_per_node 8 -m training.main \
     --save-frequency 100 \
@@ -35,8 +94,8 @@ torchrun --nproc_per_node 8 -m training.main \
     --workers=4 \
     --freeze-text \
     --use-bn-sync  \
-    --model PANN-14 \
-    --resample-method="None" \
+    --amodel PANN-14 \
+    --tmodel transformer \
     --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" \
     --datasetinfos "train" "unbalanced_train" "balanced_train"
 
@@ -60,10 +119,10 @@ torchrun --nproc_per_node 8 -m training.main \
     --epochs=200  \
     --workers=4 \
     --use-bn-sync  \
-    --model PANN-14 \
+    --amodel PANN-14 \
+    --tmodel transformer \
     --report-to "wandb" \
     --wandb-notes "text-audio-diff-lr-1e-3-1e-4-model-pann-14" \
-    --resample-method="None" \
     --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" \
     --datasetinfos "train" "unbalanced_train" "balanced_train" \
 
@@ -83,10 +142,10 @@ torchrun --nproc_per_node 8 -m training.main \
     --epochs=200  \
     --workers=4 \
     --use-bn-sync  \
-    --model PANN-14 \
+    --amodel PANN-14 \
+    --tmodel transformer \
     --report-to "wandb" \
     --wandb-notes "text-audio-diff-lr-1e-3-1e-4-model-pann-14" \
-    --resample-method="None" \
     --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" \
     --datasetinfos "train" "unbalanced_train" "balanced_train" "eval" \
 
@@ -106,10 +165,10 @@ torchrun --nproc_per_node 8 -m training.main \
     --epochs=200  \
     --workers=4 \
     --use-bn-sync  \
-    --model HTSAT-tiny \
+    --amodel HTSAT-tiny \
+    --tmodel transformer \
     --report-to "wandb" \
     --wandb-notes "text-audio-diff-lr-1e-3-1e-4-model-htsat-tiny" \
-    --resample-method="None" \
     --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" \
     --datasetinfos "train" "unbalanced_train" "balanced_train" \
 
@@ -131,10 +190,10 @@ torchrun --nproc_per_node 8 -m training.main \
     --workers=4 \
     --use-bn-sync  \
     --freeze-text \
-    --model PANN-14 \
+    --amodel PANN-14 \
+    --tmodel transformer \
     --report-to "wandb" \
     --wandb-notes "text-audio-freeze-text-lr-1e-3-model-pann-14" \
-    --resample-method="None" \
     --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" "audioset"  \
     --datasetinfos "train" "unbalanced_train" "balanced_train"\
 
@@ -153,10 +212,10 @@ torchrun --nproc_per_node 8 -m training.main \
     --workers=4 \
     --use-bn-sync  \
     --freeze-text \
-    --model HTSAT-tiny \
+    --amodel HTSAT-tiny \
+    --tmodel transformer \
     --report-to "wandb" \
     --wandb-notes "text-audio-freeze-text-lr-1e-3-model-htsat-tiny" \
-    --resample-method="None" \
     --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" "audioset"  \
     --datasetinfos "train" "unbalanced_train" "balanced_train"\
 
@@ -177,10 +236,10 @@ torchrun --nproc_per_node 8 -m training.main \
     --epochs=200  \
     --workers=4 \
     --use-bn-sync  \
-    --model PANN-14 \
+    --amodel PANN-14 \
+    --tmodel transformer \
     --report-to "wandb" \
     --wandb-notes "text-audio-diff-lr-1e-3-1e-5-freeze-after-5-model-pann-14" \
-    --resample-method="None" \
     --datasetnames "Clotho" "audiocaps" "BBCSoundEffects"  \
     --datasetinfos "train" "unbalanced_train" "balanced_train" \
 
@@ -199,10 +258,10 @@ torchrun --nproc_per_node 8 -m training.main \
     --workers=10 \
     --use-bn-sync  \
     --freeze-text \
-    --model HTSAT-tiny \
+    --amodel HTSAT-tiny \
+    --tmodel transformer \
     --report-to "wandb" \
     --wandb-notes "text-audio-freeze-text-lr-1e-3-8-dataset-model-htsat-tiny" \
-    --resample-method="None" \
     --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" "audioset" "free_to_use_sounds" "paramount_motion" "sonniss_game_effects" "wesoundeffects" \
     --datasetinfos "train" "unbalanced_train" "balanced_train"\
     --remotedata
@@ -222,10 +281,10 @@ torchrun --nproc_per_node 8 -m training.main \
     --workers=10 \
     --use-bn-sync  \
     --freeze-text \
-    --model HTSAT-tiny \
+    --amodel HTSAT-tiny \
+    --tmodel transformer \
     --report-to "wandb" \
     --wandb-notes "text-audio-freeze-text-lr-1e-3-8-dataset-model-htsat-tiny" \
-    --resample-method="None" \
     --datasetnames "audioset" \
     --datasetinfos "train" "unbalanced_train" "balanced_train"\
     --remotedata
@@ -248,10 +307,10 @@ torchrun --nproc_per_node 8  -m training.main \
     --workers=10 \
     --use-bn-sync \
     --freeze-text \
-    --model HTSAT-tiny \
+    --amodel HTSAT-tiny \
+    --tmodel transformer \
     --report-to "wandb" \
     --wandb-notes "text-audio-freeze-text-lr-1e-3-8-dataset-model-htsat-tiny" \
-    --resample-method="None" \
     --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" "audioset" "free_to_use_sounds" "paramount_motion" "sonniss_game_effects" "wesoundeffects" \
     --datasetinfos "train" "unbalanced_train" "balanced_train" "valid" \
     --top-k-checkpoint-select-dataset="Clotho-test" \
@@ -273,13 +332,102 @@ torchrun --nproc_per_node 8  -m training.main \
     --workers=10 \
     --use-bn-sync \
     --freeze-text \
-    --model PANN-14 \
+    --amodel PANN-14 \
+    --tmodel transformer \
+    --report-to "wandb" \
+    --wandb-notes "text-audio-freeze-text-lr-1e-3-8-dataset-model-pann-14" \
+    --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" "audioset" "free_to_use_sounds" "paramount_motion" "sonniss_game_effects" "wesoundeffects" \
+    --datasetinfos "train" "unbalanced_train" "balanced_train" "valid" \
+    --top-k-checkpoint-select-dataset="Clotho-test" \
+    --top-k-checkpoint-select-metric="mAP@10" \
+    --seed 3407 \
+    --remotedata
+
+torchrun --nproc_per_node 8  -m training.main \
+    --save-frequency 50 \
+    --save-top-performance 3 \
+    --save-most-recent \
+    --dataset-type="webdataset" \
+    --precision="fp32" \
+    --warmup 10000 \
+    --batch-size=184 \
+    --lr=1e-3 \
+    --wd=0.1 \
+    --epochs=400 \
+    --workers=10 \
+    --use-bn-sync \
+    --freeze-text \
+    --amodel PANN-14 \
+    --tmodel bert \
     --report-to "wandb" \
     --wandb-notes "text-audio-freeze-text-lr-1e-3-8-dataset-model-pann-14" \
-    --resample-method="None" \
     --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" "audioset" "free_to_use_sounds" "paramount_motion" "sonniss_game_effects" "wesoundeffects" \
     --datasetinfos "train" "unbalanced_train" "balanced_train" "valid" \
     --top-k-checkpoint-select-dataset="Clotho-test" \
     --top-k-checkpoint-select-metric="mAP@10" \
     --seed 3407 \
-    --remotedata
\ No newline at end of file
+    --remotedata
+
+
+torchrun --nproc_per_node 8  -m training.main \
+    --save-frequency 50 \
+    --save-top-performance 3 \
+    --save-most-recent \
+    --dataset-type="webdataset" \
+    --precision="fp32" \
+    --warmup 10000 \
+    --batch-size=184 \
+    --lr=1e-3 \
+    --wd=0.1 \
+    --epochs=400 \
+    --workers=10 \
+    --use-bn-sync \
+    --freeze-text \
+    --amodel PANN-14 \
+    --tmodel roberta \
+    --report-to "wandb" \
+    --wandb-notes "text-audio-freeze-text-lr-1e-3-8-dataset-model-pann-14" \
+    --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" "audioset" "free_to_use_sounds" "paramount_motion" "sonniss_game_effects" "wesoundeffects" \
+    --datasetinfos "train" "unbalanced_train" "balanced_train" "valid" \
+    --top-k-checkpoint-select-dataset="Clotho-test" \
+    --top-k-checkpoint-select-metric="mAP@10" \
+    --seed 3407 \
+    --remotedata
+----run on summit----
+
+module load open-ce/1.5.2-py39-0
+conda activate ~/audio_clip
+cd ~/CLAP/src
+
+num_nodes=16
+jsrun --launch_distribution=packed -bpacked:7 -n${num_nodes} -r1 -a6 -g6 -c42 ./launch-smt4.sh ./run_summit_clap.sh
+
+num_nodes=2
+jsrun --launch_distribution=packed -bpacked:7 -n${num_nodes} -r1 -a6 -g6 -c42 ./launch-smt4.sh ./run_summit_clap.sh
+
+jsrun --launch_distribution=packed -bpacked:7 -n1 -r1 -a1 -g1 -c7 ./launch-smt1.sh ./run_summit_clap_eval.sh
+
+num_nodes=1
+jsrun --launch_distribution=packed -bpacked:7 -n${num_nodes} -r1 -a6 -g6 -c42 ./launch-smt4.sh ./run_summit_clap.sh
+
+num_nodes=2
+jsrun --launch_distribution=packed -bpacked:7 -n${num_nodes} -r1 -a6 -g6 -c42 ./launch-smt4.sh ./run_summit_clap.sh
+
+jsrun --launch_distribution=packed -bpacked:7 -n1 -r1 -a1 -g1 -c7 ./launch-smt1.sh ./run_summit_clap_eval.sh
+
+num_nodes=1
+jsrun --launch_distribution=packed -bpacked:7 -n${num_nodes} -r1 -a6 -g6 -c42 ./launch-smt4.sh ./run_summit_clap.sh
+
+
+# pretrained_audio = ''
+# audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
+# audio_ckpt = audio_ckpt['state_dict']
+# save_ckpt = {}
+# keys = list(audio_ckpt.keys())
+# for key in keys:
+#     if 'audio_branch' in key:
+#         v = audio_ckpt.pop(key)
+#         save_key = key.replace('module.clap_model.', '')
+#         save_ckpt[save_key] = v
+#
+# torch.save(save_ckpt, '10.13_pann_epoch_50_mixup_processed.pt')
\ No newline at end of file
diff --git a/shs/laion-bart.sh b/shs/laion-bart.sh
new file mode 100644
index 0000000..4186b07
--- /dev/null
+++ b/shs/laion-bart.sh
@@ -0,0 +1,80 @@
+#!/bin/bash
+#SBATCH --comment clap
+#SBATCH --partition=gpu
+#SBATCH --job-name=mclap
+#SBATCH --nodes 8
+#SBATCH --ntasks-per-node 8
+#SBATCH --cpus-per-gpu=6
+#SBATCH --gres=gpu:8
+#SBATCH --output=%x_%j.out
+#SBATCH --exclude gpu-st-p4d-24xlarge-[23,30,31,108,115,134,135,183,185,186,187,188,275,277,374]
+
+module load intelmpi
+source /opt/intel/mpi/latest/env/vars.sh
+export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-install/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+export NCCL_PROTO=simple
+export PATH=/opt/amazon/efa/bin:$PATH
+export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
+
+export FI_EFA_FORK_SAFE=1
+export FI_LOG_LEVEL=1
+export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
+
+#export NCCL_ALGO=ring
+export NCCL_DEBUG=info
+#export NCCL_DEBUG_SUBSYS=INIT,ENV,GRAPH,COLL
+
+export PYTHONFAULTHANDLER=1
+
+export CUDA_LAUNCH_BLOCKING=0
+export OMPI_MCA_mtl_base_verbose=1
+export FI_EFA_ENABLE_SHM_TRANSFER=0
+export FI_PROVIDER=efa
+export FI_EFA_TX_MIN_CREDITS=64
+export NCCL_TREE_THRESHOLD=0
+
+
+#export NCCL_P2P_DISABLE=1
+#export NCCL_IBEXT_DISABLE=1
+#export NCCL_SOCKET_IFNAME="eth0,en,eth,em,bond"
+
+# sent to sub script
+export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_PORT=12802
+export COUNT_NODE=`scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l`
+
+echo go $COUNT_NODE
+echo $HOSTNAMES
+
+source /fsx/tyz/clap/bin/activate
+cd /fsx/tyz/CLAP/src
+
+export TRANSFORMERS_CACHE=/fsx/tyz/transformers_cache
+
+srun --comment clap --cpu_bind=v --accel-bind=gn python -m training.main \
+    --save-frequency 50 \
+    --save-top-performance 3 \
+    --save-most-recent \
+    --dataset-type="webdataset" \
+    --precision="fp32" \
+    --warmup 10000 \
+    --batch-size=184 \
+    --lr=1e-3 \
+    --wd=0.1 \
+    --epochs=400 \
+    --workers=10 \
+    --use-bn-sync \
+    --freeze-text \
+    --amodel PANN-14 \
+    --tmodel bart \
+    --report-to "wandb" \
+    --wandb-notes "text-audio-freeze-text-lr-1e-3-8-dataset-model-pann-14" \
+    --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" "audioset" "free_to_use_sounds" "paramount_motion" "sonniss_game_effects" "wesoundeffects" \
+    --datasetinfos "train" "unbalanced_train" "balanced_train" "valid" \
+    --top-k-checkpoint-select-dataset="Clotho-test" \
+    --top-k-checkpoint-select-metric="mAP@10" \
+    --seed 3407 \
+    --openai-model-cache-dir /fsx/tyz/transformers_cache \
+    --remotedata
+
diff --git a/shs/laion-bert.sh b/shs/laion-bert.sh
new file mode 100644
index 0000000..6a2e4cb
--- /dev/null
+++ b/shs/laion-bert.sh
@@ -0,0 +1,80 @@
+#!/bin/bash
+#SBATCH --comment clap
+#SBATCH --partition=gpu
+#SBATCH --job-name=mclap
+#SBATCH --nodes 8
+#SBATCH --ntasks-per-node 8
+#SBATCH --cpus-per-gpu=6
+#SBATCH --gres=gpu:8
+#SBATCH --output=%x_%j.out
+#SBATCH --exclude gpu-st-p4d-24xlarge-[23,30,31,108,115,134,135,183,185,186,187,188,275,277,374]
+
+module load intelmpi
+source /opt/intel/mpi/latest/env/vars.sh
+export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-install/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+export NCCL_PROTO=simple
+export PATH=/opt/amazon/efa/bin:$PATH
+export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
+
+export FI_EFA_FORK_SAFE=1
+export FI_LOG_LEVEL=1
+export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
+
+#export NCCL_ALGO=ring
+export NCCL_DEBUG=info
+#export NCCL_DEBUG_SUBSYS=INIT,ENV,GRAPH,COLL
+
+export PYTHONFAULTHANDLER=1
+
+export CUDA_LAUNCH_BLOCKING=0
+export OMPI_MCA_mtl_base_verbose=1
+export FI_EFA_ENABLE_SHM_TRANSFER=0
+export FI_PROVIDER=efa
+export FI_EFA_TX_MIN_CREDITS=64
+export NCCL_TREE_THRESHOLD=0
+
+
+#export NCCL_P2P_DISABLE=1
+#export NCCL_IBEXT_DISABLE=1
+#export NCCL_SOCKET_IFNAME="eth0,en,eth,em,bond"
+
+# sent to sub script
+export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_PORT=12802
+export COUNT_NODE=`scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l`
+
+echo go $COUNT_NODE
+echo $HOSTNAMES
+
+source /fsx/tyz/clap/bin/activate
+cd /fsx/tyz/CLAP/src
+
+export TRANSFORMERS_CACHE=/fsx/tyz/transformers_cache
+
+srun --comment clap --cpu_bind=v --accel-bind=gn python -m training.main \
+    --save-frequency 50 \
+    --save-top-performance 3 \
+    --save-most-recent \
+    --dataset-type="webdataset" \
+    --precision="fp32" \
+    --warmup 10000 \
+    --batch-size=184 \
+    --lr=1e-3 \
+    --wd=0.1 \
+    --epochs=400 \
+    --workers=10 \
+    --use-bn-sync \
+    --freeze-text \
+    --amodel PANN-14 \
+    --tmodel bert \
+    --report-to "wandb" \
+    --wandb-notes "text-audio-freeze-text-lr-1e-3-8-dataset-model-pann-14" \
+    --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" "audioset" "free_to_use_sounds" "paramount_motion" "sonniss_game_effects" "wesoundeffects" \
+    --datasetinfos "train" "unbalanced_train" "balanced_train" "valid" \
+    --top-k-checkpoint-select-dataset="Clotho-test" \
+    --top-k-checkpoint-select-metric="mAP@10" \
+    --seed 3407 \
+    --openai-model-cache-dir /fsx/tyz/transformers_cache \
+    --remotedata
+
diff --git a/shs/laion-roberta.sh b/shs/laion-roberta.sh
new file mode 100644
index 0000000..9b33626
--- /dev/null
+++ b/shs/laion-roberta.sh
@@ -0,0 +1,80 @@
+#!/bin/bash
+#SBATCH --comment clap
+#SBATCH --partition=gpu
+#SBATCH --job-name=mclap
+#SBATCH --nodes 8
+#SBATCH --ntasks-per-node 8
+#SBATCH --cpus-per-gpu=6
+#SBATCH --gres=gpu:8
+#SBATCH --output=%x_%j.out
+#SBATCH --exclude gpu-st-p4d-24xlarge-[23,30,31,108,115,134,135,183,185,186,187,188,275,277,374]
+
+module load intelmpi
+source /opt/intel/mpi/latest/env/vars.sh
+export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-install/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+export NCCL_PROTO=simple
+export PATH=/opt/amazon/efa/bin:$PATH
+export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
+
+export FI_EFA_FORK_SAFE=1
+export FI_LOG_LEVEL=1
+export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
+
+#export NCCL_ALGO=ring
+export NCCL_DEBUG=info
+#export NCCL_DEBUG_SUBSYS=INIT,ENV,GRAPH,COLL
+
+export PYTHONFAULTHANDLER=1
+
+export CUDA_LAUNCH_BLOCKING=0
+export OMPI_MCA_mtl_base_verbose=1
+export FI_EFA_ENABLE_SHM_TRANSFER=0
+export FI_PROVIDER=efa
+export FI_EFA_TX_MIN_CREDITS=64
+export NCCL_TREE_THRESHOLD=0
+
+
+#export NCCL_P2P_DISABLE=1
+#export NCCL_IBEXT_DISABLE=1
+#export NCCL_SOCKET_IFNAME="eth0,en,eth,em,bond"
+
+# sent to sub script
+export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_PORT=12802
+export COUNT_NODE=`scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l`
+
+echo go $COUNT_NODE
+echo $HOSTNAMES
+
+source /fsx/tyz/clap/bin/activate
+cd /fsx/tyz/CLAP/src
+
+export TRANSFORMERS_CACHE=/fsx/tyz/transformers_cache
+
+srun --comment clap --cpu_bind=v --accel-bind=gn python -m training.main \
+    --save-frequency 50 \
+    --save-top-performance 3 \
+    --save-most-recent \
+    --dataset-type="webdataset" \
+    --precision="fp32" \
+    --warmup 10000 \
+    --batch-size=184 \
+    --lr=1e-3 \
+    --wd=0.1 \
+    --epochs=400 \
+    --workers=10 \
+    --use-bn-sync \
+    --freeze-text \
+    --amodel PANN-14 \
+    --tmodel roberta \
+    --report-to "wandb" \
+    --wandb-notes "text-audio-freeze-text-lr-1e-3-8-dataset-model-pann-14" \
+    --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" "audioset" "free_to_use_sounds" "paramount_motion" "sonniss_game_effects" "wesoundeffects" \
+    --datasetinfos "train" "unbalanced_train" "balanced_train" "valid" \
+    --top-k-checkpoint-select-dataset="Clotho-test" \
+    --top-k-checkpoint-select-metric="mAP@10" \
+    --seed 3407 \
+    --openai-model-cache-dir /fsx/tyz/transformers_cache \
+    --remotedata
+
diff --git a/shs/laion-transformer.sh b/shs/laion-transformer.sh
new file mode 100644
index 0000000..189bb2c
--- /dev/null
+++ b/shs/laion-transformer.sh
@@ -0,0 +1,80 @@
+#!/bin/bash
+#SBATCH --comment clap
+#SBATCH --partition=gpu
+#SBATCH --job-name=mclap
+#SBATCH --nodes 8
+#SBATCH --ntasks-per-node 8
+#SBATCH --cpus-per-gpu=6
+#SBATCH --gres=gpu:8
+#SBATCH --output=%x_%j.out
+#SBATCH --exclude gpu-st-p4d-24xlarge-[23,30,31,108,115,134,135,183,185,186,187,188,275,277,374]
+
+module load intelmpi
+source /opt/intel/mpi/latest/env/vars.sh
+export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-install/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+export NCCL_PROTO=simple
+export PATH=/opt/amazon/efa/bin:$PATH
+export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
+
+export FI_EFA_FORK_SAFE=1
+export FI_LOG_LEVEL=1
+export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
+
+#export NCCL_ALGO=ring
+export NCCL_DEBUG=info
+#export NCCL_DEBUG_SUBSYS=INIT,ENV,GRAPH,COLL
+
+export PYTHONFAULTHANDLER=1
+
+export CUDA_LAUNCH_BLOCKING=0
+export OMPI_MCA_mtl_base_verbose=1
+export FI_EFA_ENABLE_SHM_TRANSFER=0
+export FI_PROVIDER=efa
+export FI_EFA_TX_MIN_CREDITS=64
+export NCCL_TREE_THRESHOLD=0
+
+
+#export NCCL_P2P_DISABLE=1
+#export NCCL_IBEXT_DISABLE=1
+#export NCCL_SOCKET_IFNAME="eth0,en,eth,em,bond"
+
+# sent to sub script
+export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_PORT=12802
+export COUNT_NODE=`scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l`
+
+echo go $COUNT_NODE
+echo $HOSTNAMES
+
+source /fsx/tyz/clap/bin/activate
+cd /fsx/tyz/CLAP/src
+
+export TRANSFORMERS_CACHE=/fsx/tyz/transformers_cache
+
+srun --comment clap --cpu_bind=v --accel-bind=gn python -m training.main \
+    --save-frequency 50 \
+    --save-top-performance 3 \
+    --save-most-recent \
+    --dataset-type="webdataset" \
+    --precision="fp32" \
+    --warmup 10000 \
+    --batch-size=184 \
+    --lr=1e-3 \
+    --wd=0.1 \
+    --epochs=400 \
+    --workers=10 \
+    --use-bn-sync \
+    --freeze-text \
+    --amodel PANN-14 \
+    --tmodel transformer \
+    --report-to "wandb" \
+    --wandb-notes "text-audio-freeze-text-lr-1e-3-8-dataset-model-pann-14" \
+    --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" "audioset" "free_to_use_sounds" "paramount_motion" "sonniss_game_effects" "wesoundeffects" \
+    --datasetinfos "train" "unbalanced_train" "balanced_train" "valid" \
+    --top-k-checkpoint-select-dataset="Clotho-test" \
+    --top-k-checkpoint-select-metric="mAP@10" \
+    --seed 3407 \
+    --openai-model-cache-dir /fsx/tyz/transformers_cache \
+    --remotedata
+
diff --git a/shs/launch-smt1.sh b/shs/launch-smt1.sh
new file mode 100755
index 0000000..730562f
--- /dev/null
+++ b/shs/launch-smt1.sh
@@ -0,0 +1,58 @@
+#!/bin/bash
+
+#
+# the following mapping is tied to "smt1"
+# this mapping may not be optimal for a given application
+#
+
+APP=$1
+
+grank=$PMIX_RANK
+lrank=$(($PMIX_RANK%6))
+export PAMI_ENABLE_STRIPING=0
+
+#if [ $grank -eq 0 ]; then 
+  #export NCCL_DEBUG_FILE="/mnt/bb/$USER/nccldeug.${LSB_JOBID}/nccldebug.${grank}"
+#fi 
+
+#LD_PRELOAD=/gpfs/alpine/lrn001/proj-shared/jqyin/directconv/directconv.so
+
+
+case ${lrank} in
+[0])
+export PAMI_IBV_DEVICE_NAME=mlx5_0:1
+export OMP_PLACES={0:7:4}
+numactl --physcpubind=0,4,8,12,16,20,24 --membind=0 $APP 
+#${APP}
+  ;;
+[1])
+export PAMI_IBV_DEVICE_NAME=mlx5_1:1
+export OMP_PLACES={28:7:4}
+numactl --physcpubind=28,32,36,40,44,48,52 --membind=0 $APP 
+#${APP}
+  ;;
+[2])
+export PAMI_IBV_DEVICE_NAME=mlx5_0:1
+export OMP_PLACES={56:7:4}
+numactl --physcpubind=56,60,64,68,72,76,80 --membind=0 $APP 
+#${APP}
+  ;;
+[3])
+export PAMI_IBV_DEVICE_NAME=mlx5_3:1
+export OMP_PLACES={88:7:4}
+numactl --physcpubind=88,92,96,100,104,108,112 --membind=8 $APP 
+#${APP}
+  ;;
+[4])
+export PAMI_IBV_DEVICE_NAME=mlx5_2:1
+export OMP_PLACES={116:7:4}
+numactl --physcpubind=116,120,124,128,132,136,140 --membind=8 $APP 
+#${APP}
+  ;;
+[5])
+export PAMI_IBV_DEVICE_NAME=mlx5_3:1
+export OMP_PLACES={144:7:4}
+numactl --physcpubind=144,148,152,156,160,164,168 --membind=8 $APP 
+#${APP}
+  ;;
+esac
diff --git a/shs/launch-smt2.sh b/shs/launch-smt2.sh
new file mode 100755
index 0000000..1dc03e9
--- /dev/null
+++ b/shs/launch-smt2.sh
@@ -0,0 +1,58 @@
+#!/bin/bash
+
+#
+# the following mapping is tied to "smt1"
+# this mapping may not be optimal for a given application
+#
+
+APP=$1
+
+grank=$PMIX_RANK
+lrank=$(($PMIX_RANK%6))
+export PAMI_ENABLE_STRIPING=0
+
+#if [ $grank -eq 0 ]; then 
+  #export NCCL_DEBUG_FILE="/mnt/bb/$USER/nccldeug.${LSB_JOBID}/nccldebug.${grank}"
+#fi 
+
+#LD_PRELOAD=/gpfs/alpine/lrn001/proj-shared/jqyin/directconv/directconv.so
+
+
+case ${lrank} in
+[0])
+export PAMI_IBV_DEVICE_NAME=mlx5_0:1
+export OMP_PLACES={0:2}:7:4
+numactl --physcpubind=0,1,4,5,8,9,12,13,16,17,20,21,24,25 --membind=0 $APP 
+#${APP}
+  ;;
+[1])
+export PAMI_IBV_DEVICE_NAME=mlx5_1:1
+export OMP_PLACES={28:2}:7:4
+numactl --physcpubind=28,29,32,33,36,37,40,41,44,45,48,49,52,53 --membind=0 $APP 
+#${APP}
+  ;;
+[2])
+export PAMI_IBV_DEVICE_NAME=mlx5_0:1
+export OMP_PLACES={56:2}:7:4
+numactl --physcpubind=56,57,60,64,68,72,76,80 --membind=0 $APP 
+#${APP}
+  ;;
+[3])
+export PAMI_IBV_DEVICE_NAME=mlx5_3:1
+export OMP_PLACES={88:7:4}
+numactl --physcpubind=88,92,96,100,104,108,112 --membind=8 $APP 
+#${APP}
+  ;;
+[4])
+export PAMI_IBV_DEVICE_NAME=mlx5_2:1
+export OMP_PLACES={116:7:4}
+numactl --physcpubind=116,120,124,128,132,136,140 --membind=8 $APP 
+#${APP}
+  ;;
+[5])
+export PAMI_IBV_DEVICE_NAME=mlx5_3:1
+export OMP_PLACES={144:7:4}
+numactl --physcpubind=144,148,152,156,160,164,168 --membind=8 $APP 
+#${APP}
+  ;;
+esac
diff --git a/src/launch-smt4.sh b/shs/launch-smt4.sh
similarity index 100%
rename from src/launch-smt4.sh
rename to shs/launch-smt4.sh
diff --git a/src/run_summit_clap.sh b/shs/run_summit_clap.sh
similarity index 67%
rename from src/run_summit_clap.sh
rename to shs/run_summit_clap.sh
index ba14033..af38b58 100755
--- a/src/run_summit_clap.sh
+++ b/shs/run_summit_clap.sh
@@ -2,14 +2,20 @@
 LOG="/gpfs/alpine/scratch/wuyusong/csc499/clap_data/audio_clip_logs/clap_log_$(date +%Y%m%d_%H%M%S).log"
 unset CUDA_VISIBLE_DEVICES
 
-source ./bootstrap_pytorch_dist_env.sh
-
-export NUMBA_CACHE_DIR='/tmp/'
+export NUMBA_CACHE_DIR=/tmp/
 
 export OMP_NUM_THREADS=7
 
 export HOROVOD_CACHE_CAPACITY=0
 
+export WANDB_DIR=/gpfs/alpine/scratch/wuyusong/csc499/clap_data/audio_clip_logs_wandb
+export WANDB_CONFIG_DIR=/gpfs/alpine/scratch/wuyusong/csc499/clap_data/audio_clip_logs_wandb
+export WANDB_MODE=offline
+wandb offline
+# Export need to be strictly before than wandb offline
+
+source ./bootstrap_pytorch_dist_env.sh
+
 python -m training.main \
     --save-frequency 50 \
     --save-top-performance 3 \
@@ -24,18 +30,17 @@ python -m training.main \
     --epochs=400 \
     --workers=4 \
     --use-bn-sync \
-    --freeze-text \
     --model HTSAT-tiny \
-    --datasetnames "audiocaps" "BBCSoundEffects" "audioset" \
+    --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" "audioset" "free_to_use_sounds" "paramount_motion" "sonniss_game_effects" "wesoundeffects" \
     --datasetinfos "train" "unbalanced_train" "balanced_train" \
-    --resample-method="None" \
-    --report-to "tensorboard" \
+    --report-to "wandb" \
+    --wandb-notes "text-audio-freeze-text-lr-1e-3-8-dataset-model" \
     --datasetpath /gpfs/alpine/scratch/wuyusong/csc499/clap_data/webdataset_tar \
     --logs /gpfs/alpine/scratch/wuyusong/csc499/clap_data/audio_clip_logs \
     --openai-model-cache-dir /gpfs/alpine/scratch/wuyusong/csc499/clap_data/.cache/clip \
-    --horovod \
+    --gather-with-grad \
     --local-loss \
-    --gather-with-grad
+    --no-eval
 
 #2>&1 > $LOG
 #    --report-to "wandb" \
diff --git a/shs/run_summit_clap_eval.sh b/shs/run_summit_clap_eval.sh
new file mode 100755
index 0000000..6cce6a1
--- /dev/null
+++ b/shs/run_summit_clap_eval.sh
@@ -0,0 +1,25 @@
+# set log path as the date and time
+LOG="/gpfs/alpine/scratch/wuyusong/csc499/clap_data/audio_clip_logs/clap_log_$(date +%Y%m%d_%H%M%S).log"
+unset CUDA_VISIBLE_DEVICES
+
+export NUMBA_CACHE_DIR=/tmp/
+
+export OMP_NUM_THREADS=7
+
+export HOROVOD_CACHE_CAPACITY=0
+
+export WANDB_DIR=/gpfs/alpine/scratch/wuyusong/csc499/clap_data/audio_clip_logs_wandb
+export WANDB_CONFIG_DIR=/gpfs/alpine/scratch/wuyusong/csc499/clap_data/audio_clip_logs_wandb
+export WANDB_MODE=offline
+wandb offline
+# Export need to be strictly before than wandb offline
+
+python -m evaluate.eval_retrieval_main \
+    --datasetpath /gpfs/alpine/scratch/wuyusong/csc499/clap_data/webdataset_tar \
+    --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" "audioset" "free_to_use_sounds" "paramount_motion" "sonniss_game_effects" "wesoundeffects" \
+    --report-to "wandb" \
+    --dataset-type="webdataset" \
+    --batch-size=96 \
+    --workers=4 \
+    --openai-model-cache-dir /gpfs/alpine/scratch/wuyusong/csc499/clap_data/.cache/clip \
+    --resume /gpfs/alpine/scratch/wuyusong/csc499/clap_data/audio_clip_logs/2022_08_28-17_16_16-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32_fewer_checkpoints/checkpoints/epoch_latest.pt
\ No newline at end of file
diff --git a/src/check_ckpt.py b/src/check_ckpt.py
new file mode 100644
index 0000000..d034dd4
--- /dev/null
+++ b/src/check_ckpt.py
@@ -0,0 +1,802 @@
+import torch
+
+def keys_in_state_dict(ckpt, device='cpu'):
+    if device=="cpu":
+        a = torch.load(ckpt, map_location=torch.device('cpu'))["state_dict"]
+    else:
+        a = torch.load(ckpt)["state_dict"]
+    print("keys_in_state_dict", a.keys())
+
+
+def check_ckpt_diff(ckpt_a, ckpt_b, key_include=None, key_exclude=None, device='cpu', verbose=True):
+    if device=="cpu":
+        a = torch.load(ckpt_a, map_location=torch.device('cpu'))["state_dict"]
+        b = torch.load(ckpt_b, map_location=torch.device('cpu'))["state_dict"]
+    else:
+        a = torch.load(ckpt_a)["state_dict"]
+        b = torch.load(ckpt_b)["state_dict"]
+    a_sum = 0
+    b_sum = 0
+    difference_count = 0
+    for k in a.keys():
+        if key_include is not None and key_include not in k:
+            continue
+        if key_exclude is not None and key_exclude in k:
+            continue
+        if k in b.keys():
+            a_sum += torch.sum(a[k])
+            b_sum += torch.sum(b[k])
+            if verbose:
+                if torch.sum(a[k]) != torch.sum(b[k]):
+                    print(f"key {k} is different")
+                    difference_count += 1
+    print("a_sum: ", a_sum)
+    print("b_sum: ", b_sum)
+    print("diff: ", a_sum - b_sum)
+    if verbose:
+        print("difference_count: ", difference_count)
+    return bool(a_sum - b_sum)
+
+# Transformer no freeze:
+# check_ckpt_diff("/fsx/clap_logs/2022_09_11-19_37_08-model_PANN-14-lr_0.001-b_160-j_4-p_fp32/checkpoints/epoch_10.pt", "/fsx/clap_logs/2022_09_11-19_37_08-model_PANN-14-lr_0.001-b_160-j_4-p_fp32/checkpoints/epoch_100.pt", "text_branch.resblocks")
+
+check_ckpt_diff("/fsx/clap_logs/2022_09_29-23_42_40-model_PANN-14-lr_0.001-b_160-j_4-p_fp32/checkpoints/epoch_1.pt",
+                "/fsx/clap_logs/2022_09_29-23_42_40-model_PANN-14-lr_0.001-b_160-j_4-p_fp32/checkpoints/epoch_2.pt",
+                "text_branch.resblocks")
+
+# key module.text_branch.resblocks.0.attn.in_proj_weight is different
+# key module.text_branch.resblocks.0.attn.in_proj_bias is different
+# key module.text_branch.resblocks.0.attn.out_proj.weight is different
+# key module.text_branch.resblocks.0.attn.out_proj.bias is different
+# key module.text_branch.resblocks.0.ln_1.weight is different
+# key module.text_branch.resblocks.0.ln_1.bias is different
+# key module.text_branch.resblocks.0.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.0.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.0.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.0.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.0.ln_2.weight is different
+# key module.text_branch.resblocks.0.ln_2.bias is different
+# key module.text_branch.resblocks.1.attn.in_proj_weight is different
+# key module.text_branch.resblocks.1.attn.in_proj_bias is different
+# key module.text_branch.resblocks.1.attn.out_proj.weight is different
+# key module.text_branch.resblocks.1.attn.out_proj.bias is different
+# key module.text_branch.resblocks.1.ln_1.weight is different
+# key module.text_branch.resblocks.1.ln_1.bias is different
+# key module.text_branch.resblocks.1.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.1.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.1.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.1.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.1.ln_2.weight is different
+# key module.text_branch.resblocks.1.ln_2.bias is different
+# key module.text_branch.resblocks.2.attn.in_proj_weight is different
+# key module.text_branch.resblocks.2.attn.in_proj_bias is different
+# key module.text_branch.resblocks.2.attn.out_proj.weight is different
+# key module.text_branch.resblocks.2.attn.out_proj.bias is different
+# key module.text_branch.resblocks.2.ln_1.weight is different
+# key module.text_branch.resblocks.2.ln_1.bias is different
+# key module.text_branch.resblocks.2.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.2.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.2.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.2.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.2.ln_2.weight is different
+# key module.text_branch.resblocks.2.ln_2.bias is different
+# key module.text_branch.resblocks.3.attn.in_proj_weight is different
+# key module.text_branch.resblocks.3.attn.in_proj_bias is different
+# key module.text_branch.resblocks.3.attn.out_proj.weight is different
+# key module.text_branch.resblocks.3.attn.out_proj.bias is different
+# key module.text_branch.resblocks.3.ln_1.weight is different
+# key module.text_branch.resblocks.3.ln_1.bias is different
+# key module.text_branch.resblocks.3.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.3.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.3.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.3.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.3.ln_2.weight is different
+# key module.text_branch.resblocks.3.ln_2.bias is different
+# key module.text_branch.resblocks.4.attn.in_proj_weight is different
+# key module.text_branch.resblocks.4.attn.in_proj_bias is different
+# key module.text_branch.resblocks.4.attn.out_proj.weight is different
+# key module.text_branch.resblocks.4.attn.out_proj.bias is different
+# key module.text_branch.resblocks.4.ln_1.weight is different
+# key module.text_branch.resblocks.4.ln_1.bias is different
+# key module.text_branch.resblocks.4.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.4.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.4.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.4.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.4.ln_2.weight is different
+# key module.text_branch.resblocks.4.ln_2.bias is different
+# key module.text_branch.resblocks.5.attn.in_proj_weight is different
+# key module.text_branch.resblocks.5.attn.in_proj_bias is different
+# key module.text_branch.resblocks.5.attn.out_proj.weight is different
+# key module.text_branch.resblocks.5.attn.out_proj.bias is different
+# key module.text_branch.resblocks.5.ln_1.weight is different
+# key module.text_branch.resblocks.5.ln_1.bias is different
+# key module.text_branch.resblocks.5.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.5.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.5.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.5.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.5.ln_2.weight is different
+# key module.text_branch.resblocks.5.ln_2.bias is different
+# key module.text_branch.resblocks.6.attn.in_proj_weight is different
+# key module.text_branch.resblocks.6.attn.in_proj_bias is different
+# key module.text_branch.resblocks.6.attn.out_proj.weight is different
+# key module.text_branch.resblocks.6.attn.out_proj.bias is different
+# key module.text_branch.resblocks.6.ln_1.weight is different
+# key module.text_branch.resblocks.6.ln_1.bias is different
+# key module.text_branch.resblocks.6.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.6.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.6.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.6.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.6.ln_2.weight is different
+# key module.text_branch.resblocks.6.ln_2.bias is different
+# key module.text_branch.resblocks.7.attn.in_proj_weight is different
+# key module.text_branch.resblocks.7.attn.in_proj_bias is different
+# key module.text_branch.resblocks.7.attn.out_proj.weight is different
+# key module.text_branch.resblocks.7.attn.out_proj.bias is different
+# key module.text_branch.resblocks.7.ln_1.weight is different
+# key module.text_branch.resblocks.7.ln_1.bias is different
+# key module.text_branch.resblocks.7.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.7.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.7.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.7.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.7.ln_2.weight is different
+# key module.text_branch.resblocks.7.ln_2.bias is different
+# key module.text_branch.resblocks.8.attn.in_proj_weight is different
+# key module.text_branch.resblocks.8.attn.in_proj_bias is different
+# key module.text_branch.resblocks.8.attn.out_proj.weight is different
+# key module.text_branch.resblocks.8.attn.out_proj.bias is different
+# key module.text_branch.resblocks.8.ln_1.weight is different
+# key module.text_branch.resblocks.8.ln_1.bias is different
+# key module.text_branch.resblocks.8.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.8.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.8.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.8.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.8.ln_2.weight is different
+# key module.text_branch.resblocks.8.ln_2.bias is different
+# key module.text_branch.resblocks.9.attn.in_proj_weight is different
+# key module.text_branch.resblocks.9.attn.in_proj_bias is different
+# key module.text_branch.resblocks.9.attn.out_proj.weight is different
+# key module.text_branch.resblocks.9.attn.out_proj.bias is different
+# key module.text_branch.resblocks.9.ln_1.weight is different
+# key module.text_branch.resblocks.9.ln_1.bias is different
+# key module.text_branch.resblocks.9.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.9.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.9.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.9.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.9.ln_2.weight is different
+# key module.text_branch.resblocks.9.ln_2.bias is different
+# key module.text_branch.resblocks.10.attn.in_proj_weight is different
+# key module.text_branch.resblocks.10.attn.in_proj_bias is different
+# key module.text_branch.resblocks.10.attn.out_proj.weight is different
+# key module.text_branch.resblocks.10.attn.out_proj.bias is different
+# key module.text_branch.resblocks.10.ln_1.weight is different
+# key module.text_branch.resblocks.10.ln_1.bias is different
+# key module.text_branch.resblocks.10.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.10.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.10.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.10.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.10.ln_2.weight is different
+# key module.text_branch.resblocks.10.ln_2.bias is different
+# key module.text_branch.resblocks.11.attn.in_proj_weight is different
+# key module.text_branch.resblocks.11.attn.in_proj_bias is different
+# key module.text_branch.resblocks.11.attn.out_proj.weight is different
+# key module.text_branch.resblocks.11.attn.out_proj.bias is different
+# key module.text_branch.resblocks.11.ln_1.weight is different
+# key module.text_branch.resblocks.11.ln_1.bias is different
+# key module.text_branch.resblocks.11.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.11.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.11.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.11.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.11.ln_2.weight is different
+# key module.text_branch.resblocks.11.ln_2.bias is different
+# a_sum:  tensor(12113.6445)
+# b_sum:  tensor(9883.4424)
+# diff:  tensor(2230.2021)
+# True
+
+
+# Transformer freeze:
+# check_ckpt_diff("/fsx/clap_logs/2022_09_16-18_55_10-model_PANN-14-lr_0.001-b_160-j_4-p_fp32/checkpoints/epoch_10.pt", "/fsx/clap_logs/2022_09_16-18_55_10-model_PANN-14-lr_0.001-b_160-j_4-p_fp32/checkpoints/epoch_100.pt", "text_branch.resblocks")
+
+# key module.text_branch.resblocks.0.attn.in_proj_weight is different
+# key module.text_branch.resblocks.0.attn.in_proj_bias is different
+# key module.text_branch.resblocks.0.attn.out_proj.weight is different
+# key module.text_branch.resblocks.0.attn.out_proj.bias is different
+# key module.text_branch.resblocks.0.ln_1.weight is different
+# key module.text_branch.resblocks.0.ln_1.bias is different
+# key module.text_branch.resblocks.0.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.0.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.0.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.0.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.0.ln_2.weight is different
+# key module.text_branch.resblocks.0.ln_2.bias is different
+# key module.text_branch.resblocks.1.attn.in_proj_weight is different
+# key module.text_branch.resblocks.1.attn.in_proj_bias is different
+# key module.text_branch.resblocks.1.attn.out_proj.weight is different
+# key module.text_branch.resblocks.1.attn.out_proj.bias is different
+# key module.text_branch.resblocks.1.ln_1.weight is different
+# key module.text_branch.resblocks.1.ln_1.bias is different
+# key module.text_branch.resblocks.1.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.1.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.1.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.1.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.1.ln_2.weight is different
+# key module.text_branch.resblocks.1.ln_2.bias is different
+# key module.text_branch.resblocks.2.attn.in_proj_weight is different
+# key module.text_branch.resblocks.2.attn.in_proj_bias is different
+# key module.text_branch.resblocks.2.attn.out_proj.weight is different
+# key module.text_branch.resblocks.2.attn.out_proj.bias is different
+# key module.text_branch.resblocks.2.ln_1.weight is different
+# key module.text_branch.resblocks.2.ln_1.bias is different
+# key module.text_branch.resblocks.2.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.2.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.2.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.2.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.2.ln_2.weight is different
+# key module.text_branch.resblocks.2.ln_2.bias is different
+# key module.text_branch.resblocks.3.attn.in_proj_weight is different
+# key module.text_branch.resblocks.3.attn.in_proj_bias is different
+# key module.text_branch.resblocks.3.attn.out_proj.weight is different
+# key module.text_branch.resblocks.3.attn.out_proj.bias is different
+# key module.text_branch.resblocks.3.ln_1.weight is different
+# key module.text_branch.resblocks.3.ln_1.bias is different
+# key module.text_branch.resblocks.3.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.3.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.3.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.3.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.3.ln_2.weight is different
+# key module.text_branch.resblocks.3.ln_2.bias is different
+# key module.text_branch.resblocks.4.attn.in_proj_weight is different
+# key module.text_branch.resblocks.4.attn.in_proj_bias is different
+# key module.text_branch.resblocks.4.attn.out_proj.weight is different
+# key module.text_branch.resblocks.4.attn.out_proj.bias is different
+# key module.text_branch.resblocks.4.ln_1.weight is different
+# key module.text_branch.resblocks.4.ln_1.bias is different
+# key module.text_branch.resblocks.4.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.4.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.4.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.4.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.4.ln_2.weight is different
+# key module.text_branch.resblocks.4.ln_2.bias is different
+# key module.text_branch.resblocks.5.attn.in_proj_weight is different
+# key module.text_branch.resblocks.5.attn.in_proj_bias is different
+# key module.text_branch.resblocks.5.attn.out_proj.weight is different
+# key module.text_branch.resblocks.5.attn.out_proj.bias is different
+# key module.text_branch.resblocks.5.ln_1.weight is different
+# key module.text_branch.resblocks.5.ln_1.bias is different
+# key module.text_branch.resblocks.5.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.5.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.5.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.5.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.5.ln_2.weight is different
+# key module.text_branch.resblocks.5.ln_2.bias is different
+# key module.text_branch.resblocks.6.attn.in_proj_weight is different
+# key module.text_branch.resblocks.6.attn.in_proj_bias is different
+# key module.text_branch.resblocks.6.attn.out_proj.weight is different
+# key module.text_branch.resblocks.6.attn.out_proj.bias is different
+# key module.text_branch.resblocks.6.ln_1.weight is different
+# key module.text_branch.resblocks.6.ln_1.bias is different
+# key module.text_branch.resblocks.6.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.6.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.6.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.6.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.6.ln_2.weight is different
+# key module.text_branch.resblocks.6.ln_2.bias is different
+# key module.text_branch.resblocks.7.attn.in_proj_weight is different
+# key module.text_branch.resblocks.7.attn.in_proj_bias is different
+# key module.text_branch.resblocks.7.attn.out_proj.weight is different
+# key module.text_branch.resblocks.7.attn.out_proj.bias is different
+# key module.text_branch.resblocks.7.ln_1.weight is different
+# key module.text_branch.resblocks.7.ln_1.bias is different
+# key module.text_branch.resblocks.7.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.7.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.7.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.7.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.7.ln_2.weight is different
+# key module.text_branch.resblocks.7.ln_2.bias is different
+# key module.text_branch.resblocks.8.attn.in_proj_weight is different
+# key module.text_branch.resblocks.8.attn.in_proj_bias is different
+# key module.text_branch.resblocks.8.attn.out_proj.weight is different
+# key module.text_branch.resblocks.8.attn.out_proj.bias is different
+# key module.text_branch.resblocks.8.ln_1.weight is different
+# key module.text_branch.resblocks.8.ln_1.bias is different
+# key module.text_branch.resblocks.8.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.8.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.8.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.8.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.8.ln_2.weight is different
+# key module.text_branch.resblocks.8.ln_2.bias is different
+# key module.text_branch.resblocks.9.attn.in_proj_weight is different
+# key module.text_branch.resblocks.9.attn.in_proj_bias is different
+# key module.text_branch.resblocks.9.attn.out_proj.weight is different
+# key module.text_branch.resblocks.9.attn.out_proj.bias is different
+# key module.text_branch.resblocks.9.ln_1.weight is different
+# key module.text_branch.resblocks.9.ln_1.bias is different
+# key module.text_branch.resblocks.9.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.9.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.9.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.9.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.9.ln_2.weight is different
+# key module.text_branch.resblocks.9.ln_2.bias is different
+# key module.text_branch.resblocks.10.attn.in_proj_weight is different
+# key module.text_branch.resblocks.10.attn.in_proj_bias is different
+# key module.text_branch.resblocks.10.attn.out_proj.weight is different
+# key module.text_branch.resblocks.10.attn.out_proj.bias is different
+# key module.text_branch.resblocks.10.ln_1.weight is different
+# key module.text_branch.resblocks.10.ln_1.bias is different
+# key module.text_branch.resblocks.10.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.10.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.10.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.10.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.10.ln_2.weight is different
+# key module.text_branch.resblocks.10.ln_2.bias is different
+# key module.text_branch.resblocks.11.attn.in_proj_weight is different
+# key module.text_branch.resblocks.11.attn.in_proj_bias is different
+# key module.text_branch.resblocks.11.attn.out_proj.weight is different
+# key module.text_branch.resblocks.11.attn.out_proj.bias is different
+# key module.text_branch.resblocks.11.ln_1.weight is different
+# key module.text_branch.resblocks.11.ln_1.bias is different
+# key module.text_branch.resblocks.11.mlp.c_fc.weight is different
+# key module.text_branch.resblocks.11.mlp.c_fc.bias is different
+# key module.text_branch.resblocks.11.mlp.c_proj.weight is different
+# key module.text_branch.resblocks.11.mlp.c_proj.bias is different
+# key module.text_branch.resblocks.11.ln_2.weight is different
+# key module.text_branch.resblocks.11.ln_2.bias is different
+# a_sum:  tensor(12133.6348)
+# b_sum:  tensor(10423.9521)
+# diff:  tensor(1709.6826)
+# True
+
+
+# bert no freeze:
+# check_ckpt_diff("/fsx/clap_logs/2022_09_14-02_33_11-model_PANN-14-lr_0.0001-b_160-j_4-p_fp32/checkpoints/epoch_10.pt", "/fsx/clap_logs/2022_09_14-02_33_11-model_PANN-14-lr_0.0001-b_160-j_4-p_fp32/checkpoints/epoch_100.pt", "text_branch.encoder")
+
+# key module.text_branch.encoder.layer.0.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.0.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.0.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.0.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.0.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.0.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.0.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.0.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.0.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.0.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.0.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.0.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.0.output.dense.weight is different
+# key module.text_branch.encoder.layer.0.output.dense.bias is different
+# key module.text_branch.encoder.layer.0.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.0.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.1.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.1.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.1.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.1.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.1.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.1.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.1.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.1.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.1.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.1.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.1.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.1.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.1.output.dense.weight is different
+# key module.text_branch.encoder.layer.1.output.dense.bias is different
+# key module.text_branch.encoder.layer.1.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.1.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.2.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.2.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.2.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.2.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.2.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.2.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.2.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.2.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.2.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.2.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.2.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.2.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.2.output.dense.weight is different
+# key module.text_branch.encoder.layer.2.output.dense.bias is different
+# key module.text_branch.encoder.layer.2.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.2.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.3.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.3.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.3.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.3.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.3.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.3.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.3.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.3.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.3.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.3.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.3.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.3.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.3.output.dense.weight is different
+# key module.text_branch.encoder.layer.3.output.dense.bias is different
+# key module.text_branch.encoder.layer.3.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.3.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.4.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.4.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.4.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.4.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.4.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.4.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.4.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.4.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.4.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.4.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.4.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.4.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.4.output.dense.weight is different
+# key module.text_branch.encoder.layer.4.output.dense.bias is different
+# key module.text_branch.encoder.layer.4.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.4.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.5.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.5.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.5.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.5.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.5.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.5.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.5.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.5.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.5.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.5.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.5.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.5.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.5.output.dense.weight is different
+# key module.text_branch.encoder.layer.5.output.dense.bias is different
+# key module.text_branch.encoder.layer.5.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.5.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.6.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.6.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.6.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.6.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.6.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.6.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.6.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.6.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.6.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.6.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.6.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.6.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.6.output.dense.weight is different
+# key module.text_branch.encoder.layer.6.output.dense.bias is different
+# key module.text_branch.encoder.layer.6.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.6.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.7.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.7.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.7.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.7.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.7.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.7.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.7.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.7.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.7.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.7.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.7.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.7.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.7.output.dense.weight is different
+# key module.text_branch.encoder.layer.7.output.dense.bias is different
+# key module.text_branch.encoder.layer.7.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.7.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.8.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.8.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.8.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.8.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.8.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.8.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.8.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.8.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.8.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.8.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.8.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.8.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.8.output.dense.weight is different
+# key module.text_branch.encoder.layer.8.output.dense.bias is different
+# key module.text_branch.encoder.layer.8.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.8.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.9.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.9.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.9.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.9.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.9.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.9.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.9.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.9.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.9.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.9.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.9.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.9.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.9.output.dense.weight is different
+# key module.text_branch.encoder.layer.9.output.dense.bias is different
+# key module.text_branch.encoder.layer.9.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.9.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.10.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.10.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.10.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.10.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.10.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.10.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.10.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.10.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.10.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.10.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.10.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.10.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.10.output.dense.weight is different
+# key module.text_branch.encoder.layer.10.output.dense.bias is different
+# key module.text_branch.encoder.layer.10.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.10.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.11.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.11.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.11.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.11.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.11.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.11.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.11.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.11.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.11.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.11.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.11.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.11.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.11.output.dense.weight is different
+# key module.text_branch.encoder.layer.11.output.dense.bias is different
+# key module.text_branch.encoder.layer.11.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.11.output.LayerNorm.bias is different
+# a_sum:  tensor(15185.1230)
+# b_sum:  tensor(15576.5596)
+# diff:  tensor(-391.4365)
+# True
+
+
+# bert freeze:
+# check_ckpt_diff("/fsx/clap_logs/2022_09_13-01_25_15-model_PANN-14-lr_0.0001-b_160-j_4-p_fp32/checkpoints/epoch_10.pt", "/fsx/clap_logs/2022_09_13-01_25_15-model_PANN-14-lr_0.0001-b_160-j_4-p_fp32/checkpoints/epoch_100.pt", "text_branch.encoder")
+
+# key module.text_branch.encoder.layer.0.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.0.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.0.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.0.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.0.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.0.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.0.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.0.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.0.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.0.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.0.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.0.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.0.output.dense.weight is different
+# key module.text_branch.encoder.layer.0.output.dense.bias is different
+# key module.text_branch.encoder.layer.0.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.0.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.1.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.1.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.1.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.1.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.1.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.1.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.1.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.1.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.1.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.1.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.1.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.1.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.1.output.dense.weight is different
+# key module.text_branch.encoder.layer.1.output.dense.bias is different
+# key module.text_branch.encoder.layer.1.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.1.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.2.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.2.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.2.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.2.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.2.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.2.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.2.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.2.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.2.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.2.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.2.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.2.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.2.output.dense.weight is different
+# key module.text_branch.encoder.layer.2.output.dense.bias is different
+# key module.text_branch.encoder.layer.2.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.2.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.3.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.3.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.3.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.3.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.3.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.3.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.3.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.3.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.3.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.3.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.3.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.3.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.3.output.dense.weight is different
+# key module.text_branch.encoder.layer.3.output.dense.bias is different
+# key module.text_branch.encoder.layer.3.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.3.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.4.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.4.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.4.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.4.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.4.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.4.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.4.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.4.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.4.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.4.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.4.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.4.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.4.output.dense.weight is different
+# key module.text_branch.encoder.layer.4.output.dense.bias is different
+# key module.text_branch.encoder.layer.4.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.4.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.5.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.5.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.5.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.5.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.5.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.5.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.5.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.5.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.5.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.5.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.5.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.5.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.5.output.dense.weight is different
+# key module.text_branch.encoder.layer.5.output.dense.bias is different
+# key module.text_branch.encoder.layer.5.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.5.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.6.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.6.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.6.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.6.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.6.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.6.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.6.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.6.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.6.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.6.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.6.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.6.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.6.output.dense.weight is different
+# key module.text_branch.encoder.layer.6.output.dense.bias is different
+# key module.text_branch.encoder.layer.6.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.6.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.7.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.7.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.7.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.7.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.7.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.7.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.7.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.7.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.7.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.7.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.7.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.7.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.7.output.dense.weight is different
+# key module.text_branch.encoder.layer.7.output.dense.bias is different
+# key module.text_branch.encoder.layer.7.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.7.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.8.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.8.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.8.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.8.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.8.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.8.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.8.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.8.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.8.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.8.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.8.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.8.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.8.output.dense.weight is different
+# key module.text_branch.encoder.layer.8.output.dense.bias is different
+# key module.text_branch.encoder.layer.8.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.8.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.9.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.9.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.9.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.9.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.9.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.9.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.9.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.9.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.9.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.9.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.9.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.9.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.9.output.dense.weight is different
+# key module.text_branch.encoder.layer.9.output.dense.bias is different
+# key module.text_branch.encoder.layer.9.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.9.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.10.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.10.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.10.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.10.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.10.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.10.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.10.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.10.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.10.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.10.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.10.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.10.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.10.output.dense.weight is different
+# key module.text_branch.encoder.layer.10.output.dense.bias is different
+# key module.text_branch.encoder.layer.10.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.10.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.11.attention.self.query.weight is different
+# key module.text_branch.encoder.layer.11.attention.self.query.bias is different
+# key module.text_branch.encoder.layer.11.attention.self.key.weight is different
+# key module.text_branch.encoder.layer.11.attention.self.key.bias is different
+# key module.text_branch.encoder.layer.11.attention.self.value.weight is different
+# key module.text_branch.encoder.layer.11.attention.self.value.bias is different
+# key module.text_branch.encoder.layer.11.attention.output.dense.weight is different
+# key module.text_branch.encoder.layer.11.attention.output.dense.bias is different
+# key module.text_branch.encoder.layer.11.attention.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.11.attention.output.LayerNorm.bias is different
+# key module.text_branch.encoder.layer.11.intermediate.dense.weight is different
+# key module.text_branch.encoder.layer.11.intermediate.dense.bias is different
+# key module.text_branch.encoder.layer.11.output.dense.weight is different
+# key module.text_branch.encoder.layer.11.output.dense.bias is different
+# key module.text_branch.encoder.layer.11.output.LayerNorm.weight is different
+# key module.text_branch.encoder.layer.11.output.LayerNorm.bias is different
+# a_sum:  tensor(15078.6641)
+# b_sum:  tensor(15540.0723)
+# diff:  tensor(-461.4082)
+# True
+
+# linear_prob_text
+# check_ckpt_diff("/fsx/clap_logs/2022_09_15-02_05_29-linear_probemodel_PANN-14-lr_0.0001-b_512-j_4-p_fp32/checkpoints/pretrain_epoch_10_lp_epoch_50.pt", "/fsx/clap_logs/2022_09_15-02_05_29-linear_probemodel_PANN-14-lr_0.0001-b_512-j_4-p_fp32/checkpoints/pretrain_epoch_10_lp_epoch_100.pt", "text_branch.resblocks")
+
+# a_sum:  tensor(12111.0244)
+# b_sum:  tensor(12111.0244)
+# diff:  tensor(0.)
+
+# linear_prob_audio
+# check_ckpt_diff("/fsx/clap_logs/2022_09_15-02_05_29-linear_probemodel_PANN-14-lr_0.0001-b_512-j_4-p_fp32/checkpoints/pretrain_epoch_10_lp_epoch_50.pt", "/fsx/clap_logs/2022_09_15-02_05_29-linear_probemodel_PANN-14-lr_0.0001-b_512-j_4-p_fp32/checkpoints/pretrain_epoch_10_lp_epoch_100.pt", "clap_model")
+
+# key clap_model.audio_branch.bn0.num_batches_tracked is different
+# key clap_model.audio_branch.conv_block1.bn1.running_mean is different
+# key clap_model.audio_branch.conv_block1.bn1.running_var is different
+# key clap_model.audio_branch.conv_block1.bn1.num_batches_tracked is different
+# key clap_model.audio_branch.conv_block1.bn2.running_mean is different
+# key clap_model.audio_branch.conv_block1.bn2.running_var is different
+# key clap_model.audio_branch.conv_block1.bn2.num_batches_tracked is different
+# key clap_model.audio_branch.conv_block2.bn1.running_mean is different
+# key clap_model.audio_branch.conv_block2.bn1.running_var is different
+# key clap_model.audio_branch.conv_block2.bn1.num_batches_tracked is different
+# key clap_model.audio_branch.conv_block2.bn2.running_mean is different
+# key clap_model.audio_branch.conv_block2.bn2.running_var is different
+# key clap_model.audio_branch.conv_block2.bn2.num_batches_tracked is different
+# key clap_model.audio_branch.conv_block3.bn1.running_mean is different
+# key clap_model.audio_branch.conv_block3.bn1.running_var is different
+# key clap_model.audio_branch.conv_block3.bn1.num_batches_tracked is different
+# key clap_model.audio_branch.conv_block3.bn2.running_mean is different
+# key clap_model.audio_branch.conv_block3.bn2.running_var is different
+# key clap_model.audio_branch.conv_block3.bn2.num_batches_tracked is different
+# key clap_model.audio_branch.conv_block4.bn1.running_mean is different
+# key clap_model.audio_branch.conv_block4.bn1.running_var is different
+# key clap_model.audio_branch.conv_block4.bn1.num_batches_tracked is different
+# key clap_model.audio_branch.conv_block4.bn2.running_mean is different
+# key clap_model.audio_branch.conv_block4.bn2.running_var is different
+# key clap_model.audio_branch.conv_block4.bn2.num_batches_tracked is different
+# key clap_model.audio_branch.conv_block5.bn1.running_mean is different
+# key clap_model.audio_branch.conv_block5.bn1.running_var is different
+# key clap_model.audio_branch.conv_block5.bn1.num_batches_tracked is different
+# key clap_model.audio_branch.conv_block5.bn2.running_mean is different
+# key clap_model.audio_branch.conv_block5.bn2.running_var is different
+# key clap_model.audio_branch.conv_block5.bn2.num_batches_tracked is different
+# key clap_model.audio_branch.conv_block6.bn1.running_mean is different
+# key clap_model.audio_branch.conv_block6.bn1.running_var is different
+# key clap_model.audio_branch.conv_block6.bn1.num_batches_tracked is different
+# key clap_model.audio_branch.conv_block6.bn2.running_mean is different
+# key clap_model.audio_branch.conv_block6.bn2.running_var is different
+# key clap_model.audio_branch.conv_block6.bn2.num_batches_tracked is different
+# a_sum:  tensor(120061.5078)
+# b_sum:  tensor(122656.0469)
+# diff:  tensor(-2594.5391)
+# True
+
diff --git a/src/dev.py b/src/dev.py
index 10a332f..09e61d9 100644
--- a/src/dev.py
+++ b/src/dev.py
@@ -61,15 +61,18 @@
 import copy
 from tqdm import tqdm
 import s3fs
+import numpy as np
 import shutil
-
+import json
+from open_clip.utils import load_p
+import logging
+lp_class_label = load_p("audioset_class_labels_indices.pkl")
 
 def log_and_continue(exn):
     """Call in an exception handler to ignore any exception, isssue a warning, and continue."""
     logging.warning(f"Handling webdataset error ({repr(exn)}). Ignoring.")
     return True
 
-
 def preprocess(
     sample,
 ):
@@ -80,6 +83,10 @@ def preprocess(
     audio_ext = "flac"
     audio_data, orig_sr = sf.read(io.BytesIO(sample[audio_ext]))
     sample["waveform"] = audio_data
+    json_dict_raw = json.loads(sample["json"].decode("utf-8"))
+    sample["multi_class_label"] = np.zeros(len(lp_class_label))
+    for x in json_dict_raw["class_names"]:
+        sample["multi_class_label"][lp_class_label[x]] = 1
     return sample
 
 
@@ -124,25 +131,26 @@ def preprocess(
 
 
 # 假设有数据集dataset_all = [‘aaa’, ‘bbb’, ‘ccc’]和他们的split
-dataset_split = {
-    "audiocaps": ["train", "valid", "test"],
-    "audioset": ["balanced_train", "unbalanced_train", "eval"],
-    "BBCSoundEffects": ["train", "test"],
-    "Clotho": ["train", "test", "valid"],
-}
-for dataset_name in ["audiocaps", "audioset", "BBCSoundEffects", "Clotho"]:
-    for split in dataset_split[dataset_name]:
-        if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
-            os.makedirs(f"./json_files/{dataset_name}/{split}")
-        os.system(
-            f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
-        )
+# dataset_split = {
+#     "audiocaps": ["train", "valid", "test"],
+#     "audioset": ["balanced_train", "unbalanced_train", "eval"],
+#     "BBCSoundEffects": ["train", "test"],
+#     "Clotho": ["train", "test", "valid"],
+# }
+# for dataset_name in ["audiocaps", "audioset", "BBCSoundEffects", "Clotho"]:
+#     for split in dataset_split[dataset_name]:
+#         if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
+#             os.makedirs(f"./json_files/{dataset_name}/{split}")
+#         os.system(
+#             f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
+#         )
 
 try:
-    input_shards = [
-        f"pipe:aws s3 cp s3://s-laion-audio/webdataset_tar/audioset/unbalanced_train/{i}.tar -"
-        for i in range(0, 3734)
-    ]
+    # input_shards = [
+    #     f"pipe:aws s3 cp s3://s-laion-audio/webdataset_tar/audioset/unbalanced_train/{i}.tar -"
+    #     for i in range(0, 3734)
+    # ]
+    input_shards = ["file://D:/eval.tar"]
     # input_shards = ["/mnt/audio_clip/webdataset_tar/audioset/eval/28.tar"]
     pipeline = [wds.SimpleShardList(input_shards)]
     _SHARD_SHUFFLE_SIZE = 2000
@@ -166,7 +174,7 @@ def preprocess(
     pipeline.extend(
         [
             wds.map(preprocess),
-            wds.to_tuple("__url__", "__key__", "waveform"),
+            wds.to_tuple("__url__", "__key__", "waveform", "multi_class_label"),
             wds.batched(1),
         ]
     )
@@ -178,7 +186,7 @@ def preprocess(
         print(k)
         old_k = k
         old_batch = copy.deepcopy(batch)
-        # print(batch)
+        print(batch)
 except:
     print(old_k)
     print(old_batch)
diff --git a/src/evaluate/eval_dcase.py b/src/evaluate/eval_dcase.py
index a32ee02..910bfe7 100644
--- a/src/evaluate/eval_dcase.py
+++ b/src/evaluate/eval_dcase.py
@@ -10,17 +10,20 @@
 import librosa
 from tqdm import tqdm
 import numpy as np
+import os
+from training.params import parse_args
 
 
-def get_output_from_single_audio(audio, text, model):
-    audio_embedding = model.audio_infer(audio, hopsize=5 * 48000, key="embedding")['embedding']
-    if audio_embedding.ndim > 1:
-        audio_embedding = audio_embedding.mean(dim=0, keepdim=True)
-    else:
-        audio_embedding = audio_embedding.unsqueeze(0)
-    audio_features = model.audio_projection(audio_embedding)
+def get_output_from_single_audio(audio, text, model, device):
+
+    # audio_embedding = model.audio_infer(audio, hopsize=5 * 48000, key="embedding", device=device)['embedding']
+    # if audio_embedding.ndim > 1:
+    #     audio_embedding = audio_embedding.mean(dim=0, keepdim=True)
+    # else:
+    #     audio_embedding = audio_embedding.unsqueeze(0)
+    audio_features = model(audio, None, device)
     audio_features = F.normalize(audio_features, dim=-1)
-    text_features = model.encode_text(text)
+    text_features = model(None, text, device=device)
     text_features = F.normalize(text_features, dim=-1)
 
     # CHANGE: before normalize or after
@@ -48,159 +51,14 @@ def get_metrics(text_to_audio_logits):
 
 
 if __name__ == '__main__':
+    args = parse_args()
 
-    # top-2 transformer + top-1 CNN
-    model_path_all = [
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_2.pt',
-    ]
-
-    model_type_all = [
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'PANN-14',
-        'PANN-14',
-        'PANN-14',
-    ]
-
-    # all transformer + top-1 CNN
-    """
-    model_path_all = [
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_2.pt',
-    ]
-
-    model_type_all = [
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'PANN-14',
-        'PANN-14',
-        'PANN-14',
-    ]
-    """
-
-    # all transformers
-    """
-    model_path_all = [
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-    ]
+    model_path = args.pretrained
 
-    model_type_all = [
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-    ]
-    """
-
-    # all systems
-    """
-    model_path_all = [
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_07-19_16_07-model_HTSAT-tiny-lr_0.001-b_96-j_4-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_05-17_53_56-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_05_31-02_55_14-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-02_55_57-model_HTSAT-tiny-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_04-05_56_46-model_PANN-14-lr_0.001-b_184-j_6-p_fp32/checkpoints/epoch_top_2.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-21_01_57-model_PANN-14-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_0.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-21_01_57-model_PANN-14-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_1.pt',
-        '/mnt/audio_clip/dcase_submission_systems/2022_06_11-21_01_57-model_PANN-14-lr_0.001-b_184-j_10-p_fp32/checkpoints/epoch_top_2.pt',
-    ]
-
-    model_type_all = [
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'HTSAT-tiny',
-        'PANN-14',
-        'PANN-14',
-        'PANN-14',
-        'PANN-14',
-        'PANN-14',
-        'PANN-14',
-    ]
-    """
-
-
-    clotho_test_preprocessed_dir = "/mnt/audio_clip/dcase_submission_systems/clotho_test/"
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--device', type=str, default='cpu')
-    args = parser.parse_args()
+    clotho_test_preprocessed_dir = "/fsx/yusong/clotho_test_set/test"
 
     cudnn.benchmark = True
     cudnn.deterministic = False
-    pretrained = 'openai'
 
     audio_features_ensemble_all = []
     text_features_ensemble_all = []
@@ -209,95 +67,86 @@ def get_metrics(text_to_audio_logits):
     logit_scale_a_ensemble_all = []
     logit_scale_t_ensemble_all = []
 
-    for model_path, model_type in zip(model_path_all, model_type_all):
-        device = torch.device(args.device)
-        model, model_cfg = create_model(
-            model_type,
-            pretrained,
-            precision='fp32',
-            device=device,
-            jit=False,
-            force_quick_gelu=False,
-        )
-
-        # load model
-        checkpoint = torch.load(model_path, map_location=device)
-        if "epoch" in checkpoint:
-            # resuming a train checkpoint w/ epoch and optimizer state
-            start_epoch = checkpoint["epoch"]
-            sd = checkpoint["state_dict"]
-            if next(iter(sd.items()))[0].startswith(
-                    "module"
-            ):
-                sd = {k[len("module."):]: v for k, v in sd.items()}
-            model.load_state_dict(sd)
-        else:
-            # loading a bare (model only) checkpoint for fine-tune or evaluation
-            model.load_state_dict(checkpoint)
 
-        model.to(device)
-        model.eval()
-        for param in model.parameters():
-            param.requires_grad = False
-
-        # take every 5th file because clotho has 5 texts for 1 audio
-        test_file_list = sorted(glob.glob(f"{clotho_test_preprocessed_dir}/*.flac"))[::5]
-
-        audio_features_all = []
-        text_features_all = []
-        audio_features_mlp_all = []
-        text_features_mlp_all = []
-        logit_scale_a_all = []
-        logit_scale_t_all = []
-
-        with torch.no_grad():
-            for file_path in tqdm(test_file_list):
-                json_path = file_path.replace(".flac", ".json")
-                with open(json_path, "r") as f:
-                    json_data = json.load(f)
-                audio, sr = librosa.load(file_path, sr=48000, mono=True)
-                audio = torch.from_numpy(audio).to(device)
-                text = json_data["original_data"]["all_captions"]
-                text = tokenize(text).to(device)
-
-                audio_features, text_features, audio_features_mlp, text_features_mlp, logit_scale_a, logit_scale_t = \
-                    get_output_from_single_audio(audio, text, model)
-
-                audio_features_all.append(audio_features.detach().cpu())
-                text_features_all.append(text_features.detach().cpu())
-                audio_features_mlp_all.append(audio_features_mlp.detach().cpu())
-                text_features_mlp_all.append(text_features_mlp.detach().cpu())
-                logit_scale_a_all.append(logit_scale_a.detach().cpu())
-                logit_scale_t_all.append(logit_scale_t.detach().cpu())
-
-        audio_features_ensemble_all.append(torch.cat(audio_features_all))
-        text_features_ensemble_all.append(torch.cat(text_features_all))
-        audio_features_mlp_ensemble_all.append(torch.cat(audio_features_mlp_all))
-        text_features_mlp_ensemble_all.append(torch.cat(text_features_mlp_all))
-        logit_scale_a_ensemble_all.append(logit_scale_a_all[0])
-        logit_scale_t_ensemble_all.append(logit_scale_t_all[0])
-
-    text_to_audio_logits_ensemble_all = []
-    for i in range(len(model_path_all)):
-        audio_features = audio_features_ensemble_all[i]
-        text_features = text_features_ensemble_all[i]
-        audio_features_mlp = audio_features_mlp_ensemble_all[i]
-        text_features_mlp = text_features_mlp_ensemble_all[i]
-        logit_scale_a = logit_scale_a_ensemble_all[i]
-        logit_scale_t = logit_scale_t_ensemble_all[i]
-
-        a_logits_per_audio = (logit_scale_a * audio_features @ text_features_mlp.t()).detach().cpu()
-        a_logits_per_text = a_logits_per_audio.t().detach().cpu()
-        t_logits_per_audio = (logit_scale_t * audio_features_mlp @ text_features.t()).detach().cpu()
-        t_logits_per_text = t_logits_per_audio.t().detach().cpu()
+    device = torch.device('cuda')
+    model, clap_model_cfg = create_model(
+        args.amodel,
+        args.tmodel,
+        args.pretrained,
+        precision=args.precision,
+        device=device,
+        jit=args.torchscript,
+        force_quick_gelu=args.force_quick_gelu,
+        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
+        skip_params=False
+    )
 
-        text_to_audio_logits = (a_logits_per_text + t_logits_per_text) / 2
-        text_to_audio_logits_ensemble_all.append(text_to_audio_logits)
+    # load model
+    checkpoint = torch.load(model_path, map_location=device)
+    if "epoch" in checkpoint:
+        # resuming a train checkpoint w/ epoch and optimizer state
+        start_epoch = checkpoint["epoch"]
+        sd = checkpoint["state_dict"]
+        if next(iter(sd.items()))[0].startswith(
+                "module"
+        ):
+            sd = {k[len("module."):]: v for k, v in sd.items()}
+        model.load_state_dict(sd)
+    else:
+        # loading a bare (model only) checkpoint for fine-tune or evaluation
+        model.load_state_dict(checkpoint)
+
+    model.to(device)
+    model.eval()
+    for param in model.parameters():
+        param.requires_grad = False
+
+    # take every 5th file because clotho has 5 texts for 1 audio
+    test_file_list = sorted(glob.glob(f"{clotho_test_preprocessed_dir}/*.flac"))
+
+    audio_features_all = []
+    text_features_all = []
+    audio_features_mlp_all = []
+    text_features_mlp_all = []
+    logit_scale_a_all = []
+    logit_scale_t_all = []
+
+    with torch.no_grad():
+        for file_path in tqdm(test_file_list):
+            json_path = file_path.replace(".flac", ".json")
+            with open(json_path, "r") as f:
+                json_data = json.load(f)
+            audio, sr = librosa.load(file_path, sr=48000, mono=True)
+            audio = torch.from_numpy(audio).to(device)
+            audio = {'waveform': audio.unsqueeze(0), 'sample_rate': sr}
+            text = json_data["text"]
+
+            if args.tmodel == "transformer":
+                from open_clip import tokenize
+                text = tokenize(text)
+            else:
+                from training.data import tokenizer
+                text = tokenizer(text)  # 5 texts for each audio
+
+            audio_features, text_features, audio_features_mlp, text_features_mlp, logit_scale_a, logit_scale_t = \
+                get_output_from_single_audio(audio, text, model, device)
+
+            audio_features_all.append(audio_features.detach().cpu())
+            text_features_all.append(text_features.detach().cpu())
+            audio_features_mlp_all.append(audio_features_mlp.detach().cpu())
+            text_features_mlp_all.append(text_features_mlp.detach().cpu())
+            logit_scale_a_all.append(logit_scale_a.detach().cpu())
+            logit_scale_t_all.append(logit_scale_t.detach().cpu())
+
+    audio_features = torch.cat(audio_features_all)
+    text_features = torch.cat(text_features_all)
+    logit_scale_a = logit_scale_a_all[0]
+
+    logits_per_audio = (logit_scale_a * audio_features @ text_features.t()).detach().cpu()
+    logits_per_text = logits_per_audio.t().detach().cpu()
 
     metrics = get_metrics(
-        torch.stack(text_to_audio_logits_ensemble_all).mean(dim=0)
+        logits_per_text
     )
 
     print(metrics)
-
-    # TODO:[Tianyu]  load from DCASE contest dataset, get the retrieval file, and save csv file
diff --git a/src/evaluate/eval_linear_probe.py b/src/evaluate/eval_linear_probe.py
new file mode 100644
index 0000000..bf5a524
--- /dev/null
+++ b/src/evaluate/eval_linear_probe.py
@@ -0,0 +1,525 @@
+'''
+Evalute the linear probe performance on different checkpoints
+'''
+from cmath import cos
+from doctest import master
+from inspect import getargs
+import logging
+import os
+from pickletools import optimize
+import random
+from datetime import datetime
+import bisect
+import copy
+import numpy as np
+import torch
+import torch.backends.cudnn as cudnn
+from torch import optim
+from torch.cuda.amp import GradScaler
+import torch.distributed as dist
+import faulthandler
+import pathlib
+import argparse
+import glob
+
+try:
+    import wandb
+except ImportError:
+    wandb = None
+
+try:
+    import torch.utils.tensorboard as tensorboard
+except ImportError:
+    tensorboard = None
+
+try:
+    import horovod.torch as hvd
+except ImportError:
+    hvd = None
+
+from open_clip import create_model_and_transforms, trace_model, create_model
+from training.data import get_data
+from training.params import parse_args
+from training.distributed import is_master, init_distributed_device, world_info_from_env
+from training.logger import setup_logging
+from training.scheduler import cosine_lr
+from training.lp_main import config_lp_optimizer
+from training.lp_train import train_one_epoch, evaluate
+from open_clip.utils import get_tar_path_from_dataset_name, dataset_split
+from open_clip.utils import load_p, load_class_label
+from open_clip.linear_probe import LinearProbe
+
+def maintain_ckpts(args, startidx, all_idx_len):
+    for i in reversed(range(startidx, all_idx_len)):
+        if os.path.exists(os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt")):
+            os.rename(
+                os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt"),
+                os.path.join(args.checkpoint_path, f"epoch_top_{i+1}.pt"),
+            )
+    if os.path.exists(
+        os.path.join(args.checkpoint_path, f"epoch_top_{all_idx_len}.pt")
+    ):
+        os.remove(os.path.join(args.checkpoint_path, f"epoch_top_{all_idx_len}.pt"))
+    return
+
+
+def update_top_k_performance(
+    new_metrics_inputs, current_top_k_ckpt_metrics, args, ckpt, bignumbetter=True, pretrain_epoch=0
+):
+    """
+    Record the top-k performance of the current epoch.
+    current_top_k_metrics is a dictionary of the form: {1: top_1_ckpt_measure, 2: top_2_ckpt_measure, ...}
+    """
+    if isinstance(new_metrics_inputs, (list, tuple)):
+        new_metrics_inputs = np.mean(new_metrics_inputs)
+        return update_top_k_performance(
+            new_metrics_inputs,
+            current_top_k_ckpt_metrics,
+            args=args,
+            ckpt=ckpt,
+            bignumbetter=bignumbetter,
+            pretrain_epoch=pretrain_epoch
+        )
+    elif isinstance(new_metrics_inputs, dict):
+        new_metrics_inputs = np.mean(list(new_metrics_inputs.values()))
+        return update_top_k_performance(
+            new_metrics_inputs,
+            current_top_k_ckpt_metrics,
+            args=args,
+            ckpt=ckpt,
+            bignumbetter=bignumbetter,
+            pretrain_epoch=pretrain_epoch
+        )
+    elif isinstance(new_metrics_inputs, (float, int)):
+        update_flag = {k: False for k in current_top_k_ckpt_metrics.keys()}
+        sorted_keys = sorted(current_top_k_ckpt_metrics.keys())
+        sorted_values = sorted(
+            current_top_k_ckpt_metrics.values(), reverse=bignumbetter
+        )
+        sorted_values_ = copy.deepcopy(sorted_values)
+        sorted_values.append(new_metrics_inputs)
+        sorted_values = sorted(sorted_values, reverse=bignumbetter)
+        sorted_values = sorted_values[:-1]
+
+        if sorted_values == sorted_values_:
+            return current_top_k_ckpt_metrics, new_metrics_inputs
+        else:
+            for i in range(len(sorted_keys)):
+                if current_top_k_ckpt_metrics[sorted_keys[i]] != sorted_values[i]:
+                    current_top_k_ckpt_metrics[sorted_keys[i]] = sorted_values[i]
+                    update_flag[sorted_keys[i]] = True
+            for i in range(len(update_flag)):
+                if update_flag[i]:
+                    maintain_ckpts(args, i, len(sorted_keys))
+                    torch.save(
+                        ckpt,
+                        os.path.join(args.checkpoint_path, f"pretrain_epoch_{pretrain_epoch}_lp_epoch_top_{i}.pt"),
+                    )
+                    break
+            return current_top_k_ckpt_metrics, new_metrics_inputs
+
+
+# def updateifNone(a, b):
+#     a = b if None else a
+#     return a
+
+
+def is_pretrained_params(n):
+    return (
+        n.startswith("clap_model.transformer")
+        or n in ["clap_model.positional_embedding", "clap_model.text_projection"]
+        or n.startswith("clap_model.token_embedding")
+        or n.startswith("clap_model.ln_final")
+        or n.startswith("clap_model.logit_scale_t")
+    )
+
+
+def random_seed(seed=42, rank=0):
+    torch.manual_seed(seed + rank)
+    np.random.seed(seed + rank)
+    random.seed(seed + rank)
+
+def main():
+    args = parse_args()
+    # sanitize model name for filesystem / uri use, easier if we don't use / in name as a rule?
+    args.amodel = args.amodel.replace("/", "-")
+    
+    pretrained_ckpts = sorted(glob.glob(os.path.join(args.pretrained, "*.pt")), key=os.path.getmtime)
+
+    if args.name is None:
+        args.name = "-".join(
+            [
+                datetime.now().strftime("%Y_%m_%d-%H_%M_%S"),
+                f"linear_probe"
+                f"model_{args.amodel}",
+                f"lr_{args.lr}",
+                f"b_{args.batch_size}",
+                f"j_{args.workers}",
+                f"p_{args.precision}",
+            ]
+        )
+
+    # discover initial world args early so we can log properly
+    args.distributed = False
+    args.local_rank, args.rank, args.world_size = world_info_from_env()
+
+    if args.remotedata and is_master(args):
+        for dataset_name in args.datasetnames:
+            for split in dataset_split[dataset_name]:
+                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
+                    os.makedirs(f"./json_files/{dataset_name}/{split}")
+                os.system(
+                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
+                )
+    args.log_path = None
+    if is_master(args, local=args.log_local):
+        log_base_path = os.path.join(args.logs, args.name)
+        os.makedirs(log_base_path, exist_ok=True)
+        log_filename = f"out-{args.rank}" if args.log_local else "out.log"
+        args.log_path = os.path.join(log_base_path, log_filename)
+
+        # avoid log dir in same name:
+        postfix = 0
+        while os.path.exists(args.log_path):
+            postfix += 1
+            log_base_path_new = log_base_path+'-'+str(postfix)
+            os.makedirs(log_base_path_new, exist_ok=True)
+            log_filename = f"out-{args.rank}" if args.log_local else "out.log"
+            args.log_path = os.path.join(log_base_path_new, log_filename)
+            # print(
+            #     "Error. Experiment already exists. Use --name {} to specify a new experiment."
+            # )
+            # return -1
+
+    # Set logger
+    args.log_level = logging.DEBUG if args.debug else logging.INFO
+    setup_logging(args.log_path, args.log_level)
+
+    # fully initialize distributed device environment
+    device = init_distributed_device(args)
+
+    args.wandb = "wandb" in args.report_to or "all" in args.report_to
+    args.tensorboard = "tensorboard" in args.report_to or "all" in args.report_to
+    if is_master(args):
+        args.tensorboard_path = (
+            os.path.join(args.logs, args.name, "tensorboard")
+            if args.tensorboard
+            else ""
+        )
+        args.checkpoint_path = os.path.join(args.logs, args.name, "checkpoints")
+        for dirname in [args.tensorboard_path, args.checkpoint_path]:
+            if dirname:
+                os.makedirs(dirname, exist_ok=True)
+    else:
+        args.tensorboard_path = ""
+        args.checkpoint_path = ""
+
+    if args.copy_codebase:
+        copy_codebase(args)
+
+    assert args.precision in ["amp", "fp16", "fp32"]
+    if args.precision == "fp16":
+        logging.warning(
+            "It is recommended to use AMP mixed-precision instead of FP16. "
+            "FP16 support needs further verification and tuning, especially for train."
+        )
+
+    if args.horovod:
+        logging.info(
+            f"Running in horovod mode with multiple processes / nodes. Device: {args.device}."
+            f"Process (global: {args.rank}, local {args.local_rank}), total {args.world_size}."
+        )
+    elif args.distributed:
+        logging.info(
+            f"Running in distributed mode with multiple processes. Device: {args.device}."
+            f"Process (global: {args.rank}, local {args.local_rank}), total {args.world_size}."
+        )
+    else:
+        logging.info(f"Running with a single process. Device {args.device}.")
+
+    logging.info(f'openai cache dir: {os.path.expanduser(args.openai_model_cache_dir)}')
+
+    # determine if this worker should save logs and checkpoints. only do so if it is rank == 0
+    args.save_logs = args.logs and args.logs.lower() != "none" and is_master(args)
+    writer = None
+    if args.save_logs and args.tensorboard:
+        assert tensorboard is not None, "Please install tensorboard."
+        writer = tensorboard.SummaryWriter(args.tensorboard_path)
+
+    if args.wandb and is_master(args):
+        assert wandb is not None, "Please install wandb."
+        logging.debug("Starting wandb.")
+        # you will have to configure this for your project!
+        wandb.init(
+            project="clap",
+            notes=args.wandb_notes,
+            name=args.wandb_notes,
+            tags=[],
+            config=vars(args),
+        )
+        logging.debug("Finished loading wandb.")
+
+    for idx, f in enumerate(pretrained_ckpts):
+        logging.info(f"pretrained on {f}")
+        args.pretrained = f
+        ckpt = torch.load(f, map_location='cpu')
+        pretrain_epoch = 0
+        if 'epoch' in ckpt:
+            pretrain_epoch = ckpt['epoch']
+        # train 
+        best_metrics = lp_main(args, device, writer, pretrain_epoch, idx)
+        
+        if args.wandb and is_master(args):
+            assert wandb is not None, "Please install wandb."
+            for name, val in best_metrics.items():
+                wandb.log({f"val/summary/{name}": val, "epoch": pretrain_epoch})
+
+    if args.wandb and is_master(args):
+        wandb.finish()
+
+def update_metric(best_metric, new_metric):
+    for key in new_metric:
+        if key not in best_metric:
+            best_metric[key] = new_metric[key]
+        else:
+            best_metric[key] = max(best_metric[key], new_metric[key])
+    return best_metric
+
+def lp_main(args, device, writer, pretrain_epoch, idx):
+
+    random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed(args.seed)
+    torch.cuda.manual_seed_all(args.seed)
+    np.random.seed(args.seed)
+    args.class_index_dict = load_class_label(args.class_label_path)
+
+
+    # Create CLAP model
+    clap_model, clap_model_cfg = create_model(
+        args.amodel,
+        args.tmodel,
+        args.pretrained,
+        precision=args.precision,
+        device=device,
+        jit=args.torchscript,
+        force_quick_gelu=args.force_quick_gelu,
+        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
+        skip_params=False,
+        enable_fusion=args.enable_fusion,
+        fusion_type=args.fusion_type
+    )
+    
+    args.lp_out_ch = len(list(args.class_index_dict.keys()))
+    # Linear Probe 
+    if idx == 0:
+        logging.info(f"linear probe using mlp: {args.lp_mlp}")
+        logging.info(f"linear probe using freeze: {args.lp_freeze}")
+        logging.info(f"linear probe act layer: {args.lp_act}")
+        logging.info(f"linear probe out ch: {args.lp_out_ch}")
+        logging.info(f"linear probe learning rate (if applicable): {args.lp_lr}")
+        logging.info(f"linear probe loss func: {args.lp_loss}")
+        logging.info(f"linear probe lp_metrics: {args.lp_metrics}")
+
+    model = LinearProbe(
+        clap_model, 
+        mlp=args.lp_mlp, freeze=args.lp_freeze, 
+        in_ch=512, out_ch=args.lp_out_ch,
+        act=args.lp_act
+    ) # in_ch is fixed (i.e., 512)
+    model = model.to(device)
+    
+    if args.horovod:
+        with torch.no_grad():
+            for param in model.parameters():
+                param.set_(param.contiguous())
+
+    if args.trace:
+        model = trace_model(model, batch_size=args.batch_size, device=device)
+
+    if is_master(args) and idx == 0:
+        logging.info("Linear Probe CLAP Model:")
+        logging.info(f"{str(clap_model)}")
+        logging.info("Params:")
+        params_file = os.path.join(args.logs, args.name, "params.txt")
+        with open(params_file, "w") as f:
+            for name in sorted(vars(args)):
+                val = getattr(args, name)
+                logging.info(f"  {name}: {val}")
+                f.write(f"{name}: {val}\n")
+    
+
+    if args.distributed and not args.horovod:
+        if args.use_bn_sync:
+            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+        ddp_args = {}
+        if args.ddp_static_graph:
+            # this doesn't exist in older PyTorch, arg only added if enabled
+            ddp_args["static_graph"] = True
+        model = torch.nn.parallel.DistributedDataParallel(
+            model, device_ids=[device], find_unused_parameters=True, **ddp_args
+        )
+
+    data = get_data(args, clap_model_cfg)
+    assert len(data), "At least one train or eval dataset must be specified."
+    if args.trace:
+        assert "train" not in data, "Cannot train with traced model"
+
+    optimizer, scheduler, text_freeze_parameters = config_lp_optimizer(model, data, args)
+
+    scaler = GradScaler() if args.precision == "amp" else None
+
+    # optionally resume from a checkpoint
+    start_epoch = 0
+    if args.resume is not None:
+        if os.path.isfile(args.resume):
+            checkpoint = torch.load(args.resume, map_location=device)
+            if "epoch" in checkpoint:
+                # resuming a train checkpoint w/ epoch and optimizer state
+                start_epoch = checkpoint["epoch"]
+                sd = checkpoint["state_dict"]
+                if not args.distributed and next(iter(sd.items()))[0].startswith(
+                    "module"
+                ):
+                    sd = {k[len("module.") :]: v for k, v in sd.items()}
+                model.load_state_dict(sd)
+                if args.split_opt:
+                    if optimizer is not None:
+                        for k, o_ in optimizer.items():
+                            o_.load_state_dict(checkpoint[k + "_" + "optimizer"])
+                if optimizer is not None:
+                    optimizer.load_state_dict(checkpoint["optimizer"])
+                if scaler is not None and "scaler" in checkpoint:
+                    scaler.load_state_dict(checkpoint["scaler"])
+                logging.info(
+                    f"=> resuming checkpoint '{args.resume}' (epoch {start_epoch})"
+                )
+            else:
+                # loading a bare (model only) checkpoint for fine-tune or evaluation
+                model.load_state_dict(checkpoint)
+                logging.info(
+                    f"=> loaded checkpoint '{args.resume}' (epoch {start_epoch})"
+                )
+            if args.freeze_text:
+                print("Freeze Text!!!!")
+                for k in text_freeze_parameters:
+                    k.requires_grad = False
+        else:
+            logging.info("=> no checkpoint found at '{}'".format(args.resume))
+
+    cudnn.benchmark = True
+    cudnn.deterministic = False
+
+    if args.wandb and is_master(args):
+        args.train_sz = data["train"].dataloader.num_samples
+        if args.val_data is not None:
+            args.val_sz = data["val"].dataloader.num_samples
+        if args.debug:
+            wandb.watch(model, log="all")
+        if idx == 0:
+            wandb.save(params_file)
+
+    best_metrics = {}
+
+    if "train" not in data:
+        metric = evaluate(model, data, start_epoch, args, writer, extra_suffix="_pe@" + str(pretrain_epoch))
+        if is_master(args):
+            best_metrics = update_metric(best_metrics, metric)
+        return
+    elif start_epoch == 0 and "val" in data and not args.no_eval:
+        metric = evaluate(model, data, 0, args, writer, extra_suffix="_pe@" + str(pretrain_epoch))
+        if is_master(args):
+            best_metrics = update_metric(best_metrics, metric)
+    if args.save_top_performance:
+        current_top_k_ckpt_metrics = {
+            i: 0 for i in range(args.save_top_performance)
+        }  # initialize the top-k metric for ckpts to 0
+
+    for epoch in range(start_epoch, args.epochs):
+        # freeze the text param after (include) args.freeze_text_after, this is -1 by default
+        if epoch == args.freeze_text_after:
+            print("Text pretrained parameters are freezed since this epoch.")
+            for k in text_freeze_parameters:
+                k.requires_grad = False
+        if is_master(args):
+            logging.info(f"Start epoch {epoch}")
+
+        train_one_epoch(model, data, epoch, optimizer, scaler, scheduler, args, writer, extra_suffix="_pe@" + str(pretrain_epoch))
+        completed_epoch = epoch + 1
+
+        if any(v in data for v in ("val", "imagenet-val", "imagenet-v2")) and not args.no_eval:
+            metric = evaluate(model, data, completed_epoch, args, writer, extra_suffix="_pe@" + str(pretrain_epoch))
+            if is_master(args):
+                best_metrics = update_metric(best_metrics, metric)
+            if args.save_top_performance:
+                top_k_dataset = args.top_k_checkpoint_select_dataset
+                top_k_metric = args.top_k_checkpoint_select_metric
+                filtered_metrics = [
+                    v
+                    for k, v in metric.items()
+                    if top_k_metric in k and top_k_dataset in k
+                ]  # check all R@10 metrics (all dataset) and use it to update the ckpt
+        # Saving checkpoints.
+        if args.save_logs:
+            opt_dict = {
+                    k + "_" + "optimizer": v.state_dict() for k, v in optimizer.items()
+                }
+            checkpoint_dict = {
+                "epoch": completed_epoch,
+                "pretrain_epoch": pretrain_epoch,
+                "name": args.name,
+                "state_dict": model.state_dict(),
+            }
+            checkpoint_dict.update(opt_dict)
+            if scaler is not None:
+                checkpoint_dict["scaler"] = scaler.state_dict()
+
+            if completed_epoch == args.epochs or (
+                args.save_frequency > 0 and (completed_epoch % args.save_frequency) == 0
+            ):
+                torch.save(
+                    checkpoint_dict,
+                    os.path.join(args.checkpoint_path, f"pretrain_epoch_{pretrain_epoch}_lp_epoch_{completed_epoch}.pt"),
+                )
+            if args.save_most_recent:
+                torch.save(
+                    checkpoint_dict,
+                    os.path.join(args.checkpoint_path, f"pretrain_epoch_{pretrain_epoch}_lp_epoch_latest.pt"),
+                )
+            if args.save_top_performance and not args.no_eval:
+                update_top_k_performance(
+                    filtered_metrics,
+                    current_top_k_ckpt_metrics,
+                    args,
+                    checkpoint_dict,
+                    bignumbetter=True,
+                    pretrain_epoch=pretrain_epoch
+                )
+    del clap_model
+    return best_metrics
+
+
+def copy_codebase(args):
+    from shutil import copytree, ignore_patterns
+
+    new_code_path = os.path.join(args.logs, args.name, "code")
+    if os.path.exists(new_code_path):
+        print(
+            f"Error. Experiment already exists at {new_code_path}. Use --name to specify a new experiment."
+        )
+        return -1
+    print(f"Copying codebase to {new_code_path}")
+    current_code_path = os.path.realpath(__file__)
+    for _ in range(3):
+        current_code_path = os.path.dirname(current_code_path)
+    copytree(
+        current_code_path, new_code_path, ignore=ignore_patterns("log", "logs", "wandb")
+    )
+    print("Done copying code.")
+    return 1
+
+
+if __name__ == "__main__":
+    main()
+
+
diff --git a/src/evaluate/eval_retrieval_main.py b/src/evaluate/eval_retrieval_main.py
new file mode 100644
index 0000000..002f2eb
--- /dev/null
+++ b/src/evaluate/eval_retrieval_main.py
@@ -0,0 +1,268 @@
+import argparse
+import os.path
+import glob
+import json
+import librosa
+from tqdm import tqdm
+import random
+import numpy as np
+import logging
+import wandb
+import time
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.backends.cudnn as cudnn
+from open_clip import create_model
+from open_clip import tokenize
+from training.logger import setup_logging
+from training.data import get_data
+from training.train import evaluate
+from open_clip.utils import get_tar_path_from_dataset_name, dataset_split
+from training.params import parse_args
+
+
+def find_params_value(file, key):
+    # find value of params in params_file
+    with open(file, 'r') as f:
+        for line in f:
+            if key + ': ' in line:
+                return line.split(': ')[1].strip()
+    return None
+
+
+def evaluate_zeroshot(model, data, start_epoch, args, writer):
+    dataloader = data["val"].dataloader
+    metrics = {}
+    device = torch.device(args.device)
+    model.eval()
+    metrics.update({"epoch": start_epoch})
+
+    all_audio_features = []
+    all_class_labels = []
+    with torch.no_grad():
+        for i, batch in enumerate(dataloader):
+            audios = batch  # contains mel_spec, wavform, and longer list
+            audio_features = model(audios, None, device)
+            audio_features = F.normalize(audio_features, dim=-1)
+            all_audio_features.append(audio_features.detach().cpu())
+            all_class_labels.append(torch.argmax(batch["class_label"], 1).long())
+        all_audio_features = torch.cat(all_audio_features, dim=0)
+        all_class_labels = torch.cat(all_class_labels, dim=0)
+        metrics["num_samples"] = all_audio_features.shape[0]
+
+        # get text features
+        all_texts = ["This is a sound of " + t for t in args.class_index_dict.keys()]
+        # (yusong): a hack, can make it better
+        if args.tmodel == "transformer":
+            from open_clip import tokenize
+            all_texts = tokenize(all_texts)
+        else:
+            from training.data import tokenizer
+            all_texts = tokenizer(all_texts)
+        all_text_features = model(None, all_texts, device)
+        all_text_features = F.normalize(all_text_features, dim=-1).detach().cpu()
+
+        # compute similarity
+        logit_scale_a, logit_scale_t = model(None, None, device)
+        logit_scale_a = logit_scale_a.cpu()
+
+        logits_per_audio = (logit_scale_a * all_audio_features @ all_text_features.t()).detach().cpu()
+        logits_per_text = logits_per_audio.t().detach().cpu()
+
+        ground_truth = all_class_labels.view(-1, 1)
+        logit = logits_per_audio
+
+        ranking = torch.argsort(logit, descending=True)
+        preds = torch.where(ranking == ground_truth)[1]  # (yusong) this line is slow because it uses single thread
+        preds = preds.detach().cpu().numpy()
+        metrics[f"{args.datasetnames[0]}_mean_rank"] = preds.mean() + 1
+        metrics[f"{args.datasetnames[0]}_median_rank"] = np.floor(np.median(preds)) + 1
+        for k in [1, 5, 10]:
+            metrics[f"{args.datasetnames[0]}_R@{k}"] = np.mean(preds < k)
+        # map@10
+        metrics[f"{args.datasetnames[0]}_mAP@10"] = np.mean(np.where(preds < 10, 1 / (preds + 1), 0.0))
+
+        logging.info(
+            f"Eval Epoch: {start_epoch} "
+            + "\t".join([f"{k}: {round(v, 4):.4f}" for k, v in metrics.items()])
+        )
+
+        if args.wandb:
+            assert wandb is not None, "Please install wandb."
+            for name, val in metrics.items():
+                wandb.log({f"val/{name}": val, "epoch": start_epoch})
+
+
+if __name__ == '__main__':
+    # (yusong) repeated run might have different metric results.
+    # This is because we randomly select crop 10s for each audio.
+    args = parse_args()
+
+    if os.path.isdir(args.pretrained):
+        log_dir = os.path.dirname(args.pretrained)
+    else:
+        log_dir = os.path.dirname(os.path.dirname(args.pretrained))
+
+    args.log_level = logging.DEBUG if args.debug else logging.INFO
+    log_path = os.path.join(log_dir, 'out.log')
+    setup_logging(log_path, args.log_level)
+    params_file = os.path.join(log_dir, 'params.txt')
+
+    seed = 3407
+    random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+
+    cudnn.benchmark = True
+    cudnn.deterministic = False
+    pretrained = 'openai'
+    amodel = find_params_value(params_file, 'amodel')
+    tmodel = find_params_value(params_file, 'tmodel')
+
+    if amodel is None or tmodel is None:
+        raise ValueError('model type not found in params file')
+
+    # set up dummy values for args
+    args.parallel_eval = False
+    args.rank = 0
+    args.local_rank = 0
+    args.world_size = 1
+    args.val_frequency = 1
+    args.epochs = 1
+    args.precision = 'fp32'
+    args.save_logs = True
+    args.wandb = True
+    args.class_index_dict = None
+
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    args.device = device
+
+    if args.remotedata:
+        for dataset_name in args.datasetnames:
+            for split in dataset_split[dataset_name]:
+                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
+                    os.makedirs(f"./json_files/{dataset_name}/{split}")
+                os.system(
+                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
+                )
+
+    if args.datasetinfos is None:
+        args.datasetinfos = ["train", "unbalanced_train", "balanced_train"]
+    if args.dataset_type == "webdataset":
+        args.train_data = get_tar_path_from_dataset_name(
+            args.datasetnames,
+            args.datasetinfos,
+            islocal=not args.remotedata,
+            proportion=args.dataset_proportion,
+            dataset_path=args.datasetpath,
+        )
+        args.val_data = get_tar_path_from_dataset_name(
+            args.datasetnames,
+            ["valid", "test", "eval"],
+            islocal=not args.remotedata,
+            proportion=1,
+            dataset_path=args.datasetpath,
+        )
+    model, model_cfg = create_model(
+        amodel,
+        tmodel,
+        pretrained,
+        precision='fp32',
+        device=device,
+        jit=False,
+        force_quick_gelu=False,
+        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
+        skip_params=False,
+        enable_fusion=args.enable_fusion,
+        fusion_type=args.fusion_type
+    )  # a hack to get model_cfg
+
+    data = get_data(args, model_cfg=model_cfg)  # (yusong): hack: no model_cfg needed to get data
+
+    writer = None  # if use tensorboard, initalize writer here
+
+    if args.wandb:
+        assert wandb is not None, "Please install wandb."
+
+        # # find the line with "wandb_notes" and get the value
+        # wandb_notes = find_params_value(params_file, 'wandb_notes')
+        # if wandb_notes is None:
+        #     print(f'wandb_notes not found in params file: {params_file}, set to timestamp.')
+        #     wandb_notes = f'experiment_{time.strftime("%Y%m%d-%H%M%S")}'
+        # wandb_notes = wandb_notes + '-eval-retrieval'
+        wandb_notes = args.wandb_notes
+
+        logging.debug("Starting wandb.")
+        args.train_sz = data["train"].dataloader.num_samples
+        if args.val_data is not None:
+            args.val_sz = data["val"].dataloader.num_samples
+        # you will have to configure this for your project!
+        if args.wandb_id is not None:
+            wandb.init(
+                project="clap",
+                id=args.wandb_id,
+                resume=True
+            )
+        else:
+            wandb.init(
+                project="clap",
+                notes=wandb_notes,
+                name=wandb_notes,
+                tags=[],
+                config=vars(args),
+            )
+        logging.debug("Finished loading wandb.")
+
+    if os.path.isdir(args.pretrained):
+        all_model_checkpoints = sorted(glob.glob(os.path.join(log_dir, 'checkpoints', '*.pt')), key=os.path.getmtime)
+    else:
+        all_model_checkpoints = [args.pretrained]
+    for model_path in all_model_checkpoints:
+        args.checkpoint_path = os.path.dirname(model_path)
+        model, model_cfg = create_model(
+            amodel,
+            tmodel,
+            pretrained,
+            precision='fp32',
+            device=device,
+            jit=False,
+            force_quick_gelu=False,
+            openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
+            skip_params=False,
+            enable_fusion=args.enable_fusion,
+            fusion_type=args.fusion_type
+        )
+
+        # load model
+        checkpoint = torch.load(model_path, map_location=device)
+        if "epoch" in checkpoint:
+            # resuming a train checkpoint w/ epoch and optimizer state
+            start_epoch = checkpoint["epoch"]
+            sd = checkpoint["state_dict"]
+            if next(iter(sd.items()))[0].startswith(
+                    "module"
+            ):
+                sd = {k[len("module."):]: v for k, v in sd.items()}
+            model.load_state_dict(sd)
+            logging.info(
+                f"=> resuming checkpoint '{model_path}' (epoch {start_epoch})"
+            )
+        else:
+            # loading a bare (model only) checkpoint for fine-tune or evaluation
+            model.load_state_dict(checkpoint)
+            start_epoch = 0
+
+        model.to(device)
+        model.eval()
+        for param in model.parameters():
+            param.requires_grad = False
+
+        if args.datasetnames == ['esc50_no_overlap'] or args.datasetnames == ['VGGSound'] or \
+                args.datasetnames == ['usd8k_no_overlap']:
+            # use the same dataset for all models
+            evaluate_zeroshot(model, data, start_epoch, args, writer)
+        else:
+            evaluate(model, data, start_epoch, args, writer)
diff --git a/src/evaluate/gen_csv.py b/src/evaluate/gen_csv.py
index 4f3339f..4ecb2e3 100644
--- a/src/evaluate/gen_csv.py
+++ b/src/evaluate/gen_csv.py
@@ -14,8 +14,8 @@
 import os
 
 
-def get_output_from_single_audio(audio, text, model):
-    audio_embedding = model.audio_infer(audio, hopsize=5 * 48000, key="embedding")[
+def get_output_from_single_audio(audio, text, model, device):
+    audio_embedding = model.audio_infer(audio, hopsize=5 * 48000, key="embedding", device=device)[
         "embedding"
     ]
     if audio_embedding.ndim > 1:
@@ -38,8 +38,8 @@ def get_output_from_single_audio(audio, text, model):
     )
 
 
-def get_audio_embeddings(audio, model):
-    audio_embedding = model.audio_infer(audio, hopsize=5 * 48000, key="embedding")[
+def get_audio_embeddings(audio, model, device):
+    audio_embedding = model.audio_infer(audio, hopsize=5 * 48000, key="embedding", device=device)[
         "embedding"
     ]
     if audio_embedding.ndim > 1:
@@ -165,7 +165,7 @@ def get_text_embeddings(text, model):
                     audio_features,
                     audio_features_mlp,
                     logit_scale_a,
-                ) = get_audio_embeddings(audio, model)
+                ) = get_audio_embeddings(audio, model,device)
 
                 audio_features_all.append(audio_features)
                 audio_features_mlp_all.append(audio_features_mlp)
diff --git a/src/linear_probe_todo.txt b/src/linear_probe_todo.txt
new file mode 100644
index 0000000..cd8664e
--- /dev/null
+++ b/src/linear_probe_todo.txt
@@ -0,0 +1,8 @@
+1. change the optimizer and scheduler  # lp_main.py 463
+2. change lp_out_ch lp_in_ch # lp_main.py 271 / 279
+3. ddp step before evaluate # lp_train.py 214-218
+4. dataset change and test  # lp_train.py 85 / 236
+5. loss does not need log_softmax/sigmoid
+6. ddp eval # loss.py 183 | # lp_train 215-219
+7. model should be load in args.pretrained
+8. what feature to use
\ No newline at end of file
diff --git a/src/open_clip/__init__.py b/src/open_clip/__init__.py
index 6c09a38..96ccf3e 100644
--- a/src/open_clip/__init__.py
+++ b/src/open_clip/__init__.py
@@ -1,5 +1,5 @@
 from .factory import list_models, create_model, create_model_and_transforms, add_model_config
-from .loss import ClipLoss
+from .loss import ClipLoss, gather_features, LPLoss, lp_gather_features, LPMetrics
 from .model import CLAP, CLAPTextCfg, CLAPVisionCfg, CLAPAudioCfp, convert_weights_to_fp16, trace_model
 from .openai import load_openai_model, list_openai_models
 from .pretrained import list_pretrained, list_pretrained_tag_models, list_pretrained_model_tags,\
diff --git a/src/open_clip/bert.py b/src/open_clip/bert.py
new file mode 100644
index 0000000..005e72d
--- /dev/null
+++ b/src/open_clip/bert.py
@@ -0,0 +1,32 @@
+from transformers import BertTokenizer, BertModel
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+model = BertModel.from_pretrained("bert-base-uncased")
+text = "Replace me by any text you'd like."
+
+def bert_embeddings(text):
+    # text = "Replace me by any text you'd like."
+    encoded_input = tokenizer(text, return_tensors='pt')
+    output = model(**encoded_input)
+    return output
+    
+from transformers import RobertaTokenizer, RobertaModel
+
+tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
+model = RobertaModel.from_pretrained('roberta-base')
+text = "Replace me by any text you'd like."
+def Roberta_embeddings(text):
+    # text = "Replace me by any text you'd like."
+    encoded_input = tokenizer(text, return_tensors='pt')
+    output = model(**encoded_input)
+    return output
+
+from transformers import BartTokenizer, BartModel
+
+tokenizer = BartTokenizer.from_pretrained('facebook/bart-base')
+model = BartModel.from_pretrained('facebook/bart-base')
+text = "Replace me by any text you'd like."
+def bart_embeddings(text):
+    # text = "Replace me by any text you'd like."
+    encoded_input = tokenizer(text, return_tensors='pt')
+    output = model(**encoded_input)
+    return output
\ No newline at end of file
diff --git a/src/open_clip/factory.py b/src/open_clip/factory.py
index ce0d142..3c3b286 100644
--- a/src/open_clip/factory.py
+++ b/src/open_clip/factory.py
@@ -18,110 +18,201 @@
 
 
 def _natural_key(string_):
-    return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())]
+    return [int(s) if s.isdigit() else s for s in re.split(r"(\d+)", string_.lower())]
 
 
 def _rescan_model_configs():
     global _MODEL_CONFIGS
 
-    config_ext = ('.json',)
+    config_ext = (".json",)
     config_files = []
     for config_path in _MODEL_CONFIG_PATHS:
         if config_path.is_file() and config_path.suffix in config_ext:
             config_files.append(config_path)
         elif config_path.is_dir():
             for ext in config_ext:
-                config_files.extend(config_path.glob(f'*{ext}'))
+                config_files.extend(config_path.glob(f"*{ext}"))
 
     for cf in config_files:
-        with open(cf, 'r') as f:
+        with open(cf, "r") as f:
             model_cfg = json.load(f)
-            if all(a in model_cfg for a in ('embed_dim', 'audio_cfg', 'text_cfg')):
+            if all(a in model_cfg for a in ("embed_dim", "audio_cfg", "text_cfg")):
                 _MODEL_CONFIGS[cf.stem] = model_cfg
 
-    _MODEL_CONFIGS = {k: v for k, v in sorted(_MODEL_CONFIGS.items(), key=lambda x: _natural_key(x[0]))}
+    _MODEL_CONFIGS = {
+        k: v
+        for k, v in sorted(_MODEL_CONFIGS.items(), key=lambda x: _natural_key(x[0]))
+    }
 
 
 _rescan_model_configs()  # initial populate of model config registry
 
 
-def load_state_dict(checkpoint_path: str, map_location='cpu'):
+def load_state_dict(checkpoint_path: str, map_location="cpu", skip_params=True):
     checkpoint = torch.load(checkpoint_path, map_location=map_location)
-    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
-        state_dict = checkpoint['state_dict']
+    if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
+        state_dict = checkpoint["state_dict"]
     else:
         state_dict = checkpoint
-    if next(iter(state_dict.items()))[0].startswith('module'):
-        state_dict = {k[7:]: v for k, v in state_dict.items()}
+    if skip_params:
+        if next(iter(state_dict.items()))[0].startswith("module"):
+            state_dict = {k[7:]: v for k, v in state_dict.items()}
+    # for k in state_dict:
+    #     if k.startswith('transformer'):
+    #         v = state_dict.pop(k)
+    #         state_dict['text_branch.' + k[12:]] = v
     return state_dict
 
 
 def create_model(
-        model_name: str,
-        pretrained: str = '',
-        precision: str = 'fp32',
-        device: torch.device = torch.device('cpu'),
-        jit: bool = False,
-        force_quick_gelu: bool = False,
-        openai_model_cache_dir: str = os.path.expanduser("~/.cache/clip"),
-        # pretrained_image: bool = False,
+    amodel_name: str,
+    tmodel_name: str,
+    pretrained: str = "",
+    precision: str = "fp32",
+    device: torch.device = torch.device("cpu"),
+    jit: bool = False,
+    force_quick_gelu: bool = False,
+    openai_model_cache_dir: str = os.path.expanduser("~/.cache/clip"),
+    skip_params=True,
+    pretrained_audio: str = "",
+    pretrained_text: str = "",
+    enable_fusion: bool = False,
+    fusion_type: str = 'None'
+    # pretrained_image: bool = False,
 ):
-    model_name = model_name.replace('/', '-')  # for callers using old naming with / in ViT names
+    amodel_name = amodel_name.replace(
+        "/", "-"
+    )  # for callers using old naming with / in ViT names
+    pretrained_orig = pretrained
     pretrained = pretrained.lower()
-    if pretrained == 'openai':
-        if model_name in _MODEL_CONFIGS:
-            logging.info(f'Loading {model_name} model config.')
-            model_cfg = deepcopy(_MODEL_CONFIGS[model_name])
+    if pretrained == "openai":
+        if amodel_name in _MODEL_CONFIGS:
+            logging.info(f"Loading {amodel_name} model config.")
+            model_cfg = deepcopy(_MODEL_CONFIGS[amodel_name])
         else:
-            logging.error(f'Model config for {model_name} not found; available models {list_models()}.')
-            raise RuntimeError(f'Model config for {model_name} not found.')
+            logging.error(
+                f"Model config for {amodel_name} not found; available models {list_models()}."
+            )
+            raise RuntimeError(f"Model config for {amodel_name} not found.")
 
-        logging.info(f'Loading pretrained ViT-B-16 text encoder from OpenAI.')
+        logging.info(f"Loading pretrained ViT-B-16 text encoder from OpenAI.")
         # Hard Code in model name
-        model = load_openai_model("ViT-B-16", model_cfg, device=device, jit=jit, cache_dir=openai_model_cache_dir)
+        model_cfg["text_cfg"]["model_type"] = tmodel_name
+        model = load_openai_model(
+            "ViT-B-16",
+            model_cfg,
+            device=device,
+            jit=jit,
+            cache_dir=openai_model_cache_dir,
+            enable_fusion=enable_fusion,
+            fusion_type=fusion_type
+        )
         # See https://discuss.pytorch.org/t/valueerror-attemting-to-unscale-fp16-gradients/81372
         if precision == "amp" or precision == "fp32":
             model = model.float()
     else:
-        if model_name in _MODEL_CONFIGS:
-            logging.info(f'Loading {model_name} model config.')
-            model_cfg = deepcopy(_MODEL_CONFIGS[model_name])
+        if amodel_name in _MODEL_CONFIGS:
+            logging.info(f"Loading {amodel_name} model config.")
+            model_cfg = deepcopy(_MODEL_CONFIGS[amodel_name])
         else:
-            logging.error(f'Model config for {model_name} not found; available models {list_models()}.')
-            raise RuntimeError(f'Model config for {model_name} not found.')
+            logging.error(
+                f"Model config for {amodel_name} not found; available models {list_models()}."
+            )
+            raise RuntimeError(f"Model config for {amodel_name} not found.")
 
         if force_quick_gelu:
             # override for use of QuickGELU on non-OpenAI transformer models
             model_cfg["quick_gelu"] = True
 
-
         # if pretrained_image:
-        #     if 'timm_model_name' in model_cfg.get('vision_cfg', {}):
+        #     if 'timm_amodel_name' in model_cfg.get('vision_cfg', {}):
         #         # pretrained weight loading for timm models set via vision_cfg
         #         model_cfg['vision_cfg']['timm_model_pretrained'] = True
         #     else:
         #         assert False, 'pretrained image towers currently only supported for timm models'
-
+        model_cfg["text_cfg"]["model_type"] = tmodel_name
+        model_cfg["enable_fusion"] = enable_fusion
+        model_cfg["fusion_type"] = fusion_type
         model = CLAP(**model_cfg)
-        
+
         if pretrained:
-            checkpoint_path = ''
-            url = get_pretrained_url(model_name, pretrained)
+            checkpoint_path = ""
+            url = get_pretrained_url(amodel_name, pretrained)
             if url:
                 checkpoint_path = download_pretrained(url, root=openai_model_cache_dir)
-            elif os.path.exists(pretrained):
-                checkpoint_path = pretrained
-
+            elif os.path.exists(pretrained_orig):
+                checkpoint_path = pretrained_orig
             if checkpoint_path:
-                logging.info(f'Loading pretrained {model_name} weights ({pretrained}).')
-                model.load_state_dict(load_state_dict(checkpoint_path))
+                logging.info(f"Loading pretrained {amodel_name}-{tmodel_name} weights ({pretrained}).")
+                ckpt = load_state_dict(checkpoint_path, skip_params=True)
+                model.load_state_dict(ckpt)
+                param_names = [n for n, p in model.named_parameters()]
+                for n in param_names:
+                    print(n, "\t", "Loaded" if n in ckpt else "Unloaded")
+            else:
+                logging.warning(
+                    f"Pretrained weights ({pretrained}) not found for model {amodel_name}."
+                )
+                raise RuntimeError(
+                    f"Pretrained weights ({pretrained}) not found for model {amodel_name}."
+                )
+
+        if pretrained_audio:
+            if amodel_name.startswith('PANN'):
+                if 'Cnn14_mAP' in pretrained_audio:  # official checkpoint
+                    audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
+                    audio_ckpt = audio_ckpt['model']
+                    keys = list(audio_ckpt.keys())
+                    for key in keys:
+                        if 'spectrogram_extractor' not in key and 'logmel_extractor' not in key:
+                            v = audio_ckpt.pop(key)
+                            audio_ckpt['audio_branch.' + key] = v
+                elif os.path.basename(pretrained_audio).startswith('PANN'):  # checkpoint trained via HTSAT codebase
+                    audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
+                    audio_ckpt = audio_ckpt['state_dict']
+                    keys = list(audio_ckpt.keys())
+                    for key in keys:
+                        if key.startswith('sed_model'):
+                            v = audio_ckpt.pop(key)
+                            audio_ckpt['audio_branch.' + key[10:]] = v
+                elif os.path.basename(pretrained_audio).startswith('finetuned'):  # checkpoint trained via linear probe codebase
+                    audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
+                else:
+                    raise ValueError('Unknown audio checkpoint')
+            elif amodel_name.startswith('HTSAT'):
+                if 'HTSAT_AudioSet_Saved' in pretrained_audio:  # official checkpoint
+                    audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
+                    audio_ckpt = audio_ckpt['state_dict']
+                    keys = list(audio_ckpt.keys())
+                    for key in keys:
+                        if key.startswith('sed_model') and ('spectrogram_extractor' not in key
+                                                            and 'logmel_extractor' not in key):
+                            v = audio_ckpt.pop(key)
+                            audio_ckpt['audio_branch.' + key[10:]] = v
+                elif os.path.basename(pretrained_audio).startswith('HTSAT'):  # checkpoint trained via HTSAT codebase
+                    audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
+                    audio_ckpt = audio_ckpt['state_dict']
+                    keys = list(audio_ckpt.keys())
+                    for key in keys:
+                        if key.startswith('sed_model'):
+                            v = audio_ckpt.pop(key)
+                            audio_ckpt['audio_branch.' + key[10:]] = v
+                elif os.path.basename(pretrained_audio).startswith('finetuned'):  # checkpoint trained via linear probe codebase
+                    audio_ckpt = torch.load(pretrained_audio, map_location='cpu')
+                else:
+                    raise ValueError('Unknown audio checkpoint')
             else:
-                logging.warning(f'Pretrained weights ({pretrained}) not found for model {model_name}.')
-                raise RuntimeError(f'Pretrained weights ({pretrained}) not found for model {model_name}.')
+                raise f'this audio encoder pretrained checkpoint is not support'
 
+            model.load_state_dict(audio_ckpt, strict=False)
+            logging.info(f"Loading pretrained {amodel_name} weights ({pretrained_audio}).")
+            param_names = [n for n, p in model.named_parameters()]
+            for n in param_names:
+                print(n, "\t", "Loaded" if n in audio_ckpt else "Unloaded")
+            
         model.to(device=device)
         if precision == "fp16":
-            assert device.type != 'cpu'
+            assert device.type != "cpu"
             convert_weights_to_fp16(model)
 
         if jit:
@@ -131,16 +222,20 @@ def create_model(
 
 
 def create_model_and_transforms(
-        model_name: str,
-        pretrained: str = '',
-        precision: str = 'fp32',
-        device: torch.device = torch.device('cpu'),
-        jit: bool = False,
-        force_quick_gelu: bool = False,
-        # pretrained_image: bool = False,
+    model_name: str,
+    pretrained: str = "",
+    precision: str = "fp32",
+    device: torch.device = torch.device("cpu"),
+    jit: bool = False,
+    force_quick_gelu: bool = False,
+    # pretrained_image: bool = False,
 ):
     model = create_model(
-        model_name, pretrained, precision, device, jit,
+        model_name,
+        pretrained,
+        precision,
+        device,
+        jit,
         force_quick_gelu=force_quick_gelu,
         # pretrained_image=pretrained_image
     )
@@ -150,12 +245,12 @@ def create_model_and_transforms(
 
 
 def list_models():
-    """ enumerate available model architectures based on config files """
+    """enumerate available model architectures based on config files"""
     return list(_MODEL_CONFIGS.keys())
 
 
 def add_model_config(path):
-    """ add model config path or file and update registry """
+    """add model config path or file and update registry"""
     if not isinstance(path, Path):
         path = Path(path)
     _MODEL_CONFIG_PATHS.append(path)
diff --git a/src/open_clip/feature_fusion.py b/src/open_clip/feature_fusion.py
new file mode 100644
index 0000000..c241951
--- /dev/null
+++ b/src/open_clip/feature_fusion.py
@@ -0,0 +1,193 @@
+'''
+Feature Fusion for Varible-Length Data Processing
+AFF/iAFF is referred and modified from https://github.com/YimianDai/open-aff/blob/master/aff_pytorch/aff_net/fusion.py
+According to the paper: Yimian Dai et al, Attentional Feature Fusion, IEEE Winter Conference on Applications of Computer Vision, WACV 2021
+'''
+
+import torch
+import torch.nn as nn
+
+
+class DAF(nn.Module):
+    '''
+    直接相加 DirectAddFuse
+    '''
+
+    def __init__(self):
+        super(DAF, self).__init__()
+
+    def forward(self, x, residual):
+        return x + residual
+
+
+class iAFF(nn.Module):
+    '''
+    多特征融合 iAFF
+    '''
+
+    def __init__(self, channels=64, r=4, type='2D'):
+        super(iAFF, self).__init__()
+        inter_channels = int(channels // r)
+
+        if type == '1D':
+            # 本地注意力
+            self.local_att = nn.Sequential(
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+
+            # 全局注意力
+            self.global_att = nn.Sequential(
+                nn.AdaptiveAvgPool1d(1),
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+
+            # 第二次本地注意力
+            self.local_att2 = nn.Sequential(
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+            # 第二次全局注意力
+            self.global_att2 = nn.Sequential(
+                nn.AdaptiveAvgPool1d(1),
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+        elif type == '2D':
+            # 本地注意力
+            self.local_att = nn.Sequential(
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+
+            # 全局注意力
+            self.global_att = nn.Sequential(
+                nn.AdaptiveAvgPool2d(1),
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+
+            # 第二次本地注意力
+            self.local_att2 = nn.Sequential(
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+            # 第二次全局注意力
+            self.global_att2 = nn.Sequential(
+                nn.AdaptiveAvgPool2d(1),
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+        else:
+            raise f'the type is not supported'
+
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x, residual):
+        flag = False
+        xa = x + residual
+        if xa.size(0) == 1:
+            xa = torch.cat([xa,xa],dim=0)
+            flag = True
+        xl = self.local_att(xa)
+        xg = self.global_att(xa)
+        xlg = xl + xg
+        wei = self.sigmoid(xlg)
+        xi = x * wei + residual * (1 - wei)
+
+        xl2 = self.local_att2(xi)
+        xg2 = self.global_att(xi)
+        xlg2 = xl2 + xg2
+        wei2 = self.sigmoid(xlg2)
+        xo = x * wei2 + residual * (1 - wei2)
+        if flag:
+            xo = xo[0].unsqueeze(0)
+        return xo
+
+
+class AFF(nn.Module):
+    '''
+    多特征融合 AFF
+    '''
+
+    def __init__(self, channels=64, r=4, type='2D'):
+        super(AFF, self).__init__()
+        inter_channels = int(channels // r)
+
+        if type == '1D':
+            self.local_att = nn.Sequential(
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+            self.global_att = nn.Sequential(
+                nn.AdaptiveAvgPool1d(1),
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+        elif type == '2D':
+            self.local_att = nn.Sequential(
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+            self.global_att = nn.Sequential(
+                nn.AdaptiveAvgPool2d(1),
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+        else:
+            raise f'the type is not supported.'
+        
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x, residual):
+        flag = False
+        xa = x + residual
+        if xa.size(0) == 1:
+            xa = torch.cat([xa,xa],dim=0)
+            flag = True
+        xl = self.local_att(xa)
+        xg = self.global_att(xa)
+        xlg = xl + xg
+        wei = self.sigmoid(xlg)
+        xo = 2 * x * wei + 2 * residual * (1 - wei)
+        if flag:
+            xo = xo[0].unsqueeze(0)
+        return xo
+
diff --git a/src/open_clip/htsat.py b/src/open_clip/htsat.py
index 3ed922e..db96116 100644
--- a/src/open_clip/htsat.py
+++ b/src/open_clip/htsat.py
@@ -24,6 +24,7 @@
 from itertools import repeat
 from .utils import do_mixup, interpolate
 
+from .feature_fusion import iAFF, AFF, DAF
 
 # from PyTorch internals
 def _ntuple(n):
@@ -70,7 +71,8 @@ def forward(self, x):
 class PatchEmbed(nn.Module):
     """ 2D Image to Patch Embedding
     """
-    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True, patch_stride = 16):
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True, patch_stride = 16,
+        enable_fusion=False, fusion_type='None'):
         super().__init__()
         img_size = to_2tuple(img_size)
         patch_size = to_2tuple(patch_size)
@@ -83,17 +85,59 @@ def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_
         self.flatten = flatten
         self.in_chans = in_chans
         self.embed_dim = embed_dim
+
+        self.enable_fusion = enable_fusion
+        self.fusion_type = fusion_type
         
         padding = ((patch_size[0] - patch_stride[0]) // 2, (patch_size[1] - patch_stride[1]) // 2)
 
-        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_stride, padding=padding)
+        if (self.enable_fusion) and (self.fusion_type == 'channel_map'):
+            self.proj = nn.Conv2d(in_chans*4, embed_dim, kernel_size=patch_size, stride=patch_stride, padding=padding)
+        else:
+            self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_stride, padding=padding)
         self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
 
-    def forward(self, x):
-        B, C, H, W = x.shape
-        assert H == self.img_size[0] and W == self.img_size[1], \
-            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
-        x = self.proj(x)
+        if (self.enable_fusion) and (self.fusion_type in ['daf_2d','aff_2d','iaff_2d']):
+            self.mel_conv2d = nn.Conv2d(in_chans, embed_dim, kernel_size=(patch_size[0], patch_size[1]*3), stride=(patch_stride[0], patch_stride[1] * 3), padding=padding)
+            if self.fusion_type == 'daf_2d':
+                self.fusion_model = DAF()
+            elif self.fusion_type == 'aff_2d':
+                self.fusion_model = AFF(channels=embed_dim, type='2D')
+            elif self.fusion_type == 'iaff_2d':
+                self.fusion_model = iAFF(channels=embed_dim, type='2D')    
+    def forward(self, x, longer_idx = None):
+        if (self.enable_fusion) and (self.fusion_type in ['daf_2d','aff_2d','iaff_2d']):
+            global_x = x[:,0:1,:,:]
+            
+
+            # global processing
+            B, C, H, W = global_x.shape
+            assert H == self.img_size[0] and W == self.img_size[1], \
+                f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+            global_x = self.proj(global_x)
+            TW = global_x.size(-1)
+            if len(longer_idx) > 0:
+                # local processing
+                local_x = x[longer_idx,1:,:,:].contiguous()
+                B, C, H, W = local_x.shape
+                local_x = local_x.view(B*C,1,H,W)
+                local_x = self.mel_conv2d(local_x)
+                local_x = local_x.view(B,C,local_x.size(1),local_x.size(2),local_x.size(3))
+                local_x = local_x.permute((0,2,3,1,4)).contiguous().flatten(3)
+                TB,TC,TH,_ = local_x.size()
+                if local_x.size(-1) < TW:
+                    local_x = torch.cat([local_x, torch.zeros((TB,TC,TH,TW-local_x.size(-1)), device=global_x.device)], dim=-1)
+                else:
+                    local_x = local_x[:,:,:,:TW]
+                
+                global_x[longer_idx] = self.fusion_model(global_x[longer_idx],local_x)
+            x = global_x
+        else:
+            B, C, H, W = x.shape
+            assert H == self.img_size[0] and W == self.img_size[1], \
+                f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+            x = self.proj(x)
+        
         if self.flatten:
             x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
         x = self.norm(x)
@@ -584,7 +628,8 @@ def __init__(self, spec_size=256, patch_size=4, patch_stride=(4,4),
                  drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
                  norm_layer=nn.LayerNorm, 
                  ape=False, patch_norm=True,
-                 use_checkpoint=False, norm_before_mlp='ln', config = None, **kwargs):
+                 use_checkpoint=False, norm_before_mlp='ln', config = None, 
+                 enable_fusion = False, fusion_type = 'None', **kwargs):
         super(HTSAT_Swin_Transformer, self).__init__()
 
         self.config = config
@@ -615,6 +660,9 @@ def __init__(self, spec_size=256, patch_size=4, patch_stride=(4,4),
 
         self.use_checkpoint = use_checkpoint
 
+        self.enable_fusion = enable_fusion
+        self.fusion_type = fusion_type
+
         #  process mel-spec ; used only once
         self.freq_ratio = self.spec_size // self.config.mel_bins
         window = 'hann'
@@ -641,7 +689,9 @@ def __init__(self, spec_size=256, patch_size=4, patch_stride=(4,4),
         # split spctrogram into non-overlapping patches
         self.patch_embed = PatchEmbed(
             img_size=self.spec_size, patch_size=self.patch_size, in_chans=self.in_chans, 
-            embed_dim=self.embed_dim, norm_layer=self.norm_layer, patch_stride = patch_stride)
+            embed_dim=self.embed_dim, norm_layer=self.norm_layer, patch_stride = patch_stride,
+            enable_fusion=self.enable_fusion, fusion_type=self.fusion_type
+            )
 
         num_patches = self.patch_embed.num_patches
         patches_resolution = self.patch_embed.grid_size
@@ -688,6 +738,19 @@ def __init__(self, spec_size=256, patch_size=4, patch_stride=(4,4),
             padding = (0,1)
         )
         self.head = nn.Linear(num_classes, num_classes)
+
+        if (self.enable_fusion) and (self.fusion_type in ['daf_1d','aff_1d','iaff_1d']):
+            self.mel_conv1d = nn.Sequential(
+                nn.Conv1d(64, 64, kernel_size=5, stride=3, padding=2),
+                nn.BatchNorm1d(64)
+            )
+            if self.fusion_type == 'daf_1d':
+                self.fusion_model = DAF()
+            elif self.fusion_type == 'aff_1d':
+                self.fusion_model = AFF(channels=64, type='1D')
+            elif self.fusion_type == 'iaff_1d':
+                self.fusion_model = iAFF(channels=64, type='1D')
+                
         self.apply(self._init_weights)
 
     def _init_weights(self, m):
@@ -708,11 +771,11 @@ def no_weight_decay_keywords(self):
         return {'relative_position_bias_table'}
 
 
-    def forward_features(self, x):
+    def forward_features(self, x, longer_idx = None):
         # A deprecated optimization for using a hierarchical output from different blocks
 
         frames_num = x.shape[2]        
-        x = self.patch_embed(x)
+        x = self.patch_embed(x, longer_idx = longer_idx)
         if self.ape:
             x = x + self.absolute_pos_embed
         x = self.pos_drop(x)
@@ -800,59 +863,113 @@ def repeat_wat2img(self, x, cur_pos):
         x = x.repeat(repeats = (1,1,4,1))
         return x
 
-    def forward(self, x: torch.Tensor, mixup_lambda = None, infer_mode = False):# out_feat_keys: List[str] = None):
-        x = self.spectrogram_extractor(x)   # (batch_size, 1, time_steps, freq_bins)
-        x = self.logmel_extractor(x)    # (batch_size, 1, time_steps, mel_bins)
-
-        x = x.transpose(1, 3)
-        x = self.bn0(x)
-        x = x.transpose(1, 3)
-        if self.training:
-            x = self.spec_augmenter(x)
-        if self.training and mixup_lambda is not None:
-            x = do_mixup(x, mixup_lambda)
-        
-        if infer_mode:
-            # in infer mode. we need to handle different length audio input
-            frame_num = x.shape[2]
-            target_T = int(self.spec_size * self.freq_ratio)
-            repeat_ratio = math.floor(target_T / frame_num)
-            x = x.repeat(repeats=(1,1,repeat_ratio,1))
+    def forward(self, x: torch.Tensor, mixup_lambda = None, infer_mode = False, device=None):# out_feat_keys: List[str] = None):
+
+        if self.enable_fusion and x["longer"].sum() == 0:
+            # if no audio is longer than 10s, then randomly select one audio to be longer
+            x["longer"][torch.randint(0, x["longer"].shape[0], (1,))] = True
+
+        if not self.enable_fusion:
+            x = x["waveform"].to(device=device, non_blocking=True)
+            x = self.spectrogram_extractor(x)   # (batch_size, 1, time_steps, freq_bins)
+            x = self.logmel_extractor(x)    # (batch_size, 1, time_steps, mel_bins)
+            x = x.transpose(1, 3)
+            x = self.bn0(x)
+            x = x.transpose(1, 3)
+            if self.training:
+                x = self.spec_augmenter(x)
+
+            if self.training and mixup_lambda is not None:
+                x = do_mixup(x, mixup_lambda)
+                
             x = self.reshape_wav2img(x)
             output_dict = self.forward_features(x)
         else:
-            if x.shape[2] > self.freq_ratio * self.spec_size:
-                if self.training:
-                    x = self.crop_wav(x, crop_size=self.freq_ratio * self.spec_size)
-                    x = self.reshape_wav2img(x)
-                    output_dict = self.forward_features(x)
+            longer_list = x["longer"].to(device=device, non_blocking=True)
+            x = x["mel_fusion"].to(device=device, non_blocking=True)
+            x = x.transpose(1, 3)
+            x = self.bn0(x)
+            x = x.transpose(1, 3)
+            longer_list_idx = torch.where(longer_list)[0]
+            if self.fusion_type in ['daf_1d','aff_1d','iaff_1d']:
+                new_x = x[:,0:1,:,:].clone().contiguous()
+                if len(longer_list_idx) > 0:
+                # local processing
+                    fusion_x_local = x[longer_list_idx,1:,:,:].clone().contiguous()
+                    FB,FC,FT,FF = fusion_x_local.size()
+                    fusion_x_local = fusion_x_local.view(FB * FC, FT, FF)
+                    fusion_x_local = torch.permute(fusion_x_local, (0,2,1)).contiguous()
+                    fusion_x_local = self.mel_conv1d(fusion_x_local)
+                    fusion_x_local = fusion_x_local.view(FB,FC,FF,fusion_x_local.size(-1))
+                    fusion_x_local = torch.permute(fusion_x_local, (0,2,1,3)).contiguous().flatten(2)
+                    if fusion_x_local.size(-1) < FT:
+                        fusion_x_local = torch.cat([fusion_x_local, torch.zeros((FB,FF,FT- fusion_x_local.size(-1)), device=device)], dim=-1)
+                    else:
+                        fusion_x_local = fusion_x_local[:,:,:FT]
+                    # 1D fusion
+                    new_x = new_x.squeeze(1).permute((0,2,1)).contiguous()
+                    new_x[longer_list_idx] = self.fusion_model(new_x[longer_list_idx], fusion_x_local)
+                    x = new_x.permute((0,2,1)).contiguous()[:,None,:,:]
                 else:
-                    # Change: Hard code here
-                    overlap_size = (x.shape[2] - 1) // 4
-                    output_dicts = []
-                    crop_size = (x.shape[2] - 1) // 2
-                    for cur_pos in range(0, x.shape[2] - crop_size - 1, overlap_size):
-                        tx = self.crop_wav(x, crop_size = crop_size, spe_pos = cur_pos)
-                        tx = self.reshape_wav2img(tx)
-                        output_dicts.append(self.forward_features(tx))
-                    clipwise_output = torch.zeros_like(output_dicts[0]["clipwise_output"]).float().to(x.device)
-                    framewise_output = torch.zeros_like(output_dicts[0]["framewise_output"]).float().to(x.device)
-                    for d in output_dicts:
-                        clipwise_output += d["clipwise_output"]
-                        framewise_output += d["framewise_output"]
-                    clipwise_output  = clipwise_output / len(output_dicts)
-                    framewise_output = framewise_output / len(output_dicts)
-                    output_dict = {
-                        'framewise_output': framewise_output, 
-                        'clipwise_output': clipwise_output
-                    }
-            else: # this part is typically used, and most easy one
-                x = self.reshape_wav2img(x)
-                output_dict = self.forward_features(x)
+                    x = new_x
+
+            elif self.fusion_type in ['daf_2d','aff_2d','iaff_2d','channel_map']:
+                x = x # no change
+
+            if self.training:
+                x = self.spec_augmenter(x)
+            if self.training and mixup_lambda is not None:
+                x = do_mixup(x, mixup_lambda)
+
+            x = self.reshape_wav2img(x)
+            output_dict = self.forward_features(x, longer_idx = longer_list_idx)
+       
+        # if infer_mode:
+        #     # in infer mode. we need to handle different length audio input
+        #     frame_num = x.shape[2]
+        #     target_T = int(self.spec_size * self.freq_ratio)
+        #     repeat_ratio = math.floor(target_T / frame_num)
+        #     x = x.repeat(repeats=(1,1,repeat_ratio,1))
+        #     x = self.reshape_wav2img(x)
+        #     output_dict = self.forward_features(x)
+        # else:
+        #     if x.shape[2] > self.freq_ratio * self.spec_size:
+        #         if self.training:
+        #             x = self.crop_wav(x, crop_size=self.freq_ratio * self.spec_size)
+        #             x = self.reshape_wav2img(x)
+        #             output_dict = self.forward_features(x)
+        #         else:
+        #             # Change: Hard code here
+        #             overlap_size = (x.shape[2] - 1) // 4
+        #             output_dicts = []
+        #             crop_size = (x.shape[2] - 1) // 2
+        #             for cur_pos in range(0, x.shape[2] - crop_size - 1, overlap_size):
+        #                 tx = self.crop_wav(x, crop_size = crop_size, spe_pos = cur_pos)
+        #                 tx = self.reshape_wav2img(tx)
+        #                 output_dicts.append(self.forward_features(tx))
+        #             clipwise_output = torch.zeros_like(output_dicts[0]["clipwise_output"]).float().to(x.device)
+        #             framewise_output = torch.zeros_like(output_dicts[0]["framewise_output"]).float().to(x.device)
+        #             for d in output_dicts:
+        #                 clipwise_output += d["clipwise_output"]
+        #                 framewise_output += d["framewise_output"]
+        #             clipwise_output  = clipwise_output / len(output_dicts)
+        #             framewise_output = framewise_output / len(output_dicts)
+        #             output_dict = {
+        #                 'framewise_output': framewise_output, 
+        #                 'clipwise_output': clipwise_output
+        #             }
+        #     else: # this part is typically used, and most easy one
+        #         x = self.reshape_wav2img(x)
+        #         output_dict = self.forward_features(x)
         # x = self.head(x)
+
+        # We process the data in the dataloader part, in that here we only consider the input_T < fixed_T
+        
+        
+
         return output_dict
 
-def create_htsat_model(audio_cfg):
+def create_htsat_model(audio_cfg, enable_fusion=False, fusion_type='None'):
     try:
 
         assert audio_cfg.model_name in ["tiny", "base", "large"], "model name for HTS-AT is wrong!"
@@ -866,7 +983,9 @@ def create_htsat_model(audio_cfg):
                 depths=[2,2,6,2],
                 num_heads=[4,8,16,32],
                 window_size=8,
-                config = audio_cfg
+                config = audio_cfg,
+                enable_fusion = enable_fusion,
+                fusion_type = fusion_type
             )
         elif audio_cfg.model_name == "base":
             model = HTSAT_Swin_Transformer(
@@ -878,7 +997,9 @@ def create_htsat_model(audio_cfg):
                 depths=[2,2,12,2],
                 num_heads=[4,8,16,32],
                 window_size=8,
-                config = audio_cfg
+                config = audio_cfg,
+                enable_fusion = enable_fusion,
+                fusion_type = fusion_type
             )
         elif audio_cfg.model_name == "large":
             model = HTSAT_Swin_Transformer(
@@ -890,7 +1011,9 @@ def create_htsat_model(audio_cfg):
                 depths=[2,2,12,2],
                 num_heads=[4,8,16,32],
                 window_size=8,
-                config = audio_cfg
+                config = audio_cfg,
+                enable_fusion = enable_fusion,
+                fusion_type = fusion_type
             )
         
         return model
diff --git a/src/open_clip/linear_probe.py b/src/open_clip/linear_probe.py
new file mode 100644
index 0000000..bb2841d
--- /dev/null
+++ b/src/open_clip/linear_probe.py
@@ -0,0 +1,63 @@
+import numpy as np
+import torch.nn.functional as F
+from torch import nn
+from .model import MLPLayers
+
+
+class LinearProbe(nn.Module):
+    def __init__(self, model, mlp, freeze, in_ch, out_ch, act=None):
+        """
+        Args:
+            model: nn.Module
+            mlp: bool, if True, then use the MLP layer as the linear probe module
+            freeze: bool, if Ture, then freeze all the CLAP model's layers when training the linear probe
+            in_ch: int, the output channel from CLAP model
+            out_ch: int, the output channel from linear probe (class_num)
+            act: torch.nn.functional, the activation function before the loss function
+        """
+        super().__init__()
+        in_ch = 512
+        self.clap_model = model
+        self.clap_model.text_branch = None  # to save memory
+        self.freeze = freeze
+        if mlp:
+            self.lp_layer = MLPLayers(units=[in_ch, in_ch * 2, out_ch])
+        else:
+            self.lp_layer = nn.Linear(in_ch, out_ch)
+
+        if self.freeze:
+            for param in self.clap_model.parameters():
+                param.requires_grad = False
+
+        if act == 'None':
+            self.act = None
+        elif act == 'relu':
+            self.act = nn.ReLU()
+        elif act == 'elu':
+            self.act = nn.ELU()
+        elif act == 'prelu':
+            self.act = nn.PReLU(num_parameters=in_ch)
+        elif act == 'softmax':
+            self.act = nn.Softmax(dim=-1)
+        elif act == 'sigmoid':
+            self.act = nn.Sigmoid()
+
+    def forward(self, x, mix_lambda=None, device=None):
+        """
+        Args:
+            x: waveform, torch.tensor [batch, t_samples] / batch of mel_spec and longer list
+            mix_lambda: torch.tensor [batch], the mixup lambda
+        Returns:
+            class_prob: torch.tensor [batch, class_num]
+
+        """
+        # batchnorm cancel grandient
+        if self.freeze:
+            self.clap_model.eval()
+
+        x = self.clap_model.audio_projection(
+            self.clap_model.audio_branch(x, mixup_lambda=mix_lambda, device=device)["embedding"])
+        out = self.lp_layer(x)
+        if self.act is not None:
+            out = self.act(out)
+        return out
diff --git a/src/open_clip/loss.py b/src/open_clip/loss.py
index e75e1c8..53bbedd 100644
--- a/src/open_clip/loss.py
+++ b/src/open_clip/loss.py
@@ -1,7 +1,10 @@
+from multiprocessing.sharedctypes import Value
 import torch
 import torch.distributed.nn
 from torch import distributed as dist, nn as nn
 from torch.nn import functional as F
+import numpy as np
+from sklearn.metrics import average_precision_score, roc_auc_score, accuracy_score 
 
 try:
     import horovod.torch as hvd
@@ -12,70 +15,80 @@
 def gather_features(
         audio_features,
         text_features,
-        audio_features_mlp, 
-        text_features_mlp,
+        audio_features_mlp=None, 
+        text_features_mlp=None,
         local_loss=False,
         gather_with_grad=False,
         rank=0,
         world_size=1,
-        use_horovod=False
+        use_horovod=False,
+        mlp_loss=False
 ):
     if use_horovod:
         assert hvd is not None, 'Please install horovod'
         if gather_with_grad:
             all_audio_features = hvd.allgather(audio_features)
             all_text_features = hvd.allgather(text_features)
-            all_audio_features_mlp = hvd.allgather(audio_features_mlp)
-            all_text_features_mlp = hvd.allgather(text_features_mlp)
+            if mlp_loss:
+                all_audio_features_mlp = hvd.allgather(audio_features_mlp)
+                all_text_features_mlp = hvd.allgather(text_features_mlp)
         else:
             with torch.no_grad():
                 all_audio_features = hvd.allgather(audio_features)
                 all_text_features = hvd.allgather(text_features)
-                all_audio_features_mlp = hvd.allgather(audio_features_mlp)
-                all_text_features_mlp = hvd.allgather(text_features_mlp)
+                if mlp_loss:
+                    all_audio_features_mlp = hvd.allgather(audio_features_mlp)
+                    all_text_features_mlp = hvd.allgather(text_features_mlp)
             if not local_loss:
                 # ensure grads for local rank when all_* features don't have a gradient
                 gathered_audio_features = list(all_audio_features.chunk(world_size, dim=0))
                 gathered_text_features = list(all_text_features.chunk(world_size, dim=0))
-                gathered_audio_features_mlp = list(all_audio_features_mlp.chunk(world_size, dim=0))
-                gathered_text_features_mlp = list(all_text_features_mlp.chunk(world_size, dim=0))
                 gathered_audio_features[rank] = audio_features
                 gathered_text_features[rank] = text_features
-                gathered_audio_features_mlp[rank] = audio_features_mlp
-                gathered_text_features_mlp[rank] = text_features_mlp
                 all_audio_features = torch.cat(gathered_audio_features, dim=0)
                 all_text_features = torch.cat(gathered_text_features, dim=0)
-                all_audio_features_mlp = torch.cat(gathered_audio_features_mlp, dim=0)
-                all_text_features_mlp = torch.cat(gathered_text_features_mlp, dim=0)
+                if mlp_loss:
+                    gathered_audio_features_mlp = list(all_audio_features_mlp.chunk(world_size, dim=0))
+                    gathered_text_features_mlp = list(all_text_features_mlp.chunk(world_size, dim=0))
+                    gathered_audio_features_mlp[rank] = audio_features_mlp
+                    gathered_text_features_mlp[rank] = text_features_mlp
+                    all_audio_features_mlp = torch.cat(gathered_audio_features_mlp, dim=0)
+                    all_text_features_mlp = torch.cat(gathered_text_features_mlp, dim=0)
     else:
         # We gather tensors from all gpus
         if gather_with_grad:
             all_audio_features = torch.cat(torch.distributed.nn.all_gather(audio_features), dim=0)
             all_text_features = torch.cat(torch.distributed.nn.all_gather(text_features), dim=0)
-            all_audio_features_mlp = torch.cat(torch.distributed.nn.all_gather(audio_features_mlp), dim=0)
-            all_text_features_mlp = torch.cat(torch.distributed.nn.all_gather(text_features_mlp), dim=0)
+            if mlp_loss:
+                all_audio_features_mlp = torch.cat(torch.distributed.nn.all_gather(audio_features_mlp), dim=0)
+                all_text_features_mlp = torch.cat(torch.distributed.nn.all_gather(text_features_mlp), dim=0)
         else:
             gathered_audio_features = [torch.zeros_like(audio_features) for _ in range(world_size)]
             gathered_text_features = [torch.zeros_like(text_features) for _ in range(world_size)]
-            gathered_audio_features_mlp = [torch.zeros_like(audio_features_mlp) for _ in range(world_size)]
-            gathered_text_features_mlp = [torch.zeros_like(text_features_mlp) for _ in range(world_size)]
             dist.all_gather(gathered_audio_features, audio_features)
             dist.all_gather(gathered_text_features, text_features)
-            dist.all_gather(gathered_audio_features_mlp, audio_features_mlp)
-            dist.all_gather(gathered_text_features_mlp, text_features_mlp)
+            if mlp_loss:
+                gathered_audio_features_mlp = [torch.zeros_like(audio_features_mlp) for _ in range(world_size)]
+                gathered_text_features_mlp = [torch.zeros_like(text_features_mlp) for _ in range(world_size)]
+                dist.all_gather(gathered_audio_features_mlp, audio_features_mlp)
+                dist.all_gather(gathered_text_features_mlp, text_features_mlp)
             if not local_loss:
                 # ensure grads for local rank when all_* features don't have a gradient
                 gathered_audio_features[rank] = audio_features
                 gathered_text_features[rank] = text_features
-                gathered_audio_features_mlp[rank] = audio_features_mlp
-                gathered_text_features_mlp[rank] = text_features_mlp
+                if mlp_loss:
+                    gathered_audio_features_mlp[rank] = audio_features_mlp
+                    gathered_text_features_mlp[rank] = text_features_mlp
+
             all_audio_features = torch.cat(gathered_audio_features, dim=0)
             all_text_features = torch.cat(gathered_text_features, dim=0)
-            all_audio_features_mlp = torch.cat(gathered_audio_features_mlp, dim=0)
-            all_text_features_mlp = torch.cat(gathered_text_features_mlp, dim=0)
-
-    return all_audio_features, all_text_features, all_audio_features_mlp, all_text_features_mlp
-
+            if mlp_loss:
+                all_audio_features_mlp = torch.cat(gathered_audio_features_mlp, dim=0)
+                all_text_features_mlp = torch.cat(gathered_text_features_mlp, dim=0)
+    if mlp_loss:
+        return all_audio_features, all_text_features, all_audio_features_mlp, all_text_features_mlp
+    else:
+        return all_audio_features, all_text_features
 
 class ClipLoss(nn.Module):
 
@@ -87,6 +100,8 @@ def __init__(
             rank=0,
             world_size=1,
             use_horovod=False,
+            mlp_loss=False,
+            weight_loss_kappa=0,
     ):
         super().__init__()
         self.local_loss = local_loss
@@ -95,51 +110,198 @@ def __init__(
         self.rank = rank
         self.world_size = world_size
         self.use_horovod = use_horovod
-
+        self.mlp_loss = mlp_loss
+        self.weighted_loss = bool(weight_loss_kappa!=0)
+        self.weight_loss_kappa = weight_loss_kappa
         # cache state
         self.prev_num_logits = 0
         self.labels = {}
 
-    def forward(self, audio_features, text_features, audio_features_mlp, text_features_mlp, logit_scale_a, logit_scale_t):
+    def forward(self, audio_features, text_features, logit_scale_a, logit_scale_t=None, audio_features_mlp=None, text_features_mlp=None):
         device = audio_features.device
-        if self.world_size > 1:
-            all_audio_features, all_text_features, all_audio_features_mlp, all_text_features_mlp = gather_features(
-                audio_features, text_features, audio_features_mlp, text_features_mlp,
-                self.local_loss, self.gather_with_grad, self.rank, self.world_size, self.use_horovod)
-
-            if self.local_loss:
-                a_logits_per_audio = logit_scale_a * audio_features @ all_text_features_mlp.T
-                a_logits_per_text = logit_scale_a * text_features_mlp @ all_audio_features.T
-                t_logits_per_audio = logit_scale_t * audio_features_mlp @ all_text_features.T
-                t_logits_per_text = logit_scale_t * text_features @ all_audio_features_mlp.T
+        if self.mlp_loss:
+            if self.world_size > 1:
+                all_audio_features, all_text_features, all_audio_features_mlp, all_text_features_mlp = gather_features(
+                    audio_features=audio_features,text_features=text_features,
+                    audio_features_mlp=audio_features_mlp,text_features_mlp=text_features_mlp,
+                    local_loss=self.local_loss,gather_with_grad=self.gather_with_grad,
+                    rank=self.rank,world_size=self.world_size,use_horovod=self.use_horovod,
+                    mlp_loss=self.mlp_loss
+                )
+                if self.local_loss:
+                    a_logits_per_audio = logit_scale_a * audio_features @ all_text_features_mlp.T
+                    a_logits_per_text = logit_scale_a * text_features_mlp @ all_audio_features.T
+                    t_logits_per_audio = logit_scale_t * audio_features_mlp @ all_text_features.T
+                    t_logits_per_text = logit_scale_t * text_features @ all_audio_features_mlp.T
+                else:
+                    a_logits_per_audio = logit_scale_a * all_audio_features @ all_text_features_mlp.T
+                    a_logits_per_text = a_logits_per_audio.T
+                    t_logits_per_audio = logit_scale_t * all_audio_features_mlp @ all_text_features.T
+                    t_logits_per_text = t_logits_per_audio.T
             else:
-                a_logits_per_audio = logit_scale_a * all_audio_features @ all_text_features_mlp.T
-                a_logits_per_text = a_logits_per_audio.T
-                t_logits_per_audio = logit_scale_t * all_audio_features_mlp @ all_text_features.T
-                t_logits_per_text = t_logits_per_audio.T
-        else:
-            a_logits_per_audio = logit_scale_a * audio_features @ text_features_mlp.T
-            a_logits_per_text = logit_scale_a * text_features_mlp @ audio_features.T
-            t_logits_per_audio = logit_scale_t * audio_features_mlp @ text_features.T
-            t_logits_per_text = logit_scale_t * text_features @ audio_features_mlp.T
-
-        # calculated ground-truth and cache if enabled
-        num_logits = a_logits_per_audio.shape[0]
-        if self.prev_num_logits != num_logits or device not in self.labels:
-            labels = torch.arange(num_logits, device=device, dtype=torch.long)
-            if self.world_size > 1 and self.local_loss:
-                labels = labels + num_logits * self.rank
-            if self.cache_labels:
-                self.labels[device] = labels
-                self.prev_num_logits = num_logits
+                a_logits_per_audio = logit_scale_a * audio_features @ text_features_mlp.T
+                a_logits_per_text = logit_scale_a * text_features_mlp @ audio_features.T
+                t_logits_per_audio = logit_scale_t * audio_features_mlp @ text_features.T
+                t_logits_per_text = logit_scale_t * text_features @ audio_features_mlp.T
+
+            # calculated ground-truth and cache if enabled
+            num_logits = a_logits_per_audio.shape[0]
+            if self.prev_num_logits != num_logits or device not in self.labels:
+                labels = torch.arange(num_logits, device=device, dtype=torch.long)
+                if self.world_size > 1 and self.local_loss:
+                    labels = labels + num_logits * self.rank
+                if self.cache_labels:
+                    self.labels[device] = labels
+                    self.prev_num_logits = num_logits
+            else:
+                labels = self.labels[device]
+
+            if not self.weighted_loss:
+                total_loss = (
+                    F.cross_entropy(a_logits_per_audio, labels) +
+                    F.cross_entropy(a_logits_per_text, labels) + 
+                    F.cross_entropy(t_logits_per_audio, labels) +
+                    F.cross_entropy(t_logits_per_text, labels) 
+                    ) / 4
+            else:
+                audio_weight = (audio_features@audio_features.T).detach()
+                audio_weight = (torch.exp(torch.sum(audio_weight, axis=1)/(self.weight_loss_kappa*len(audio_weight)))).detach()
+                text_weight = (text_features@text_features.T).detach()
+                text_weight = (torch.exp(torch.sum(text_weight, axis=1)/(self.weight_loss_kappa*len(text_features)))).detach()
+                total_loss = (
+                    F.cross_entropy(a_logits_per_audio, labels, weight=audio_weight) +
+                    F.cross_entropy(a_logits_per_text, labels, weight=audio_weight) + 
+                    F.cross_entropy(t_logits_per_audio, labels, weight=text_weight) +
+                    F.cross_entropy(t_logits_per_text, labels, weight=text_weight) 
+                    ) / 4
         else:
-            labels = self.labels[device]
+            if self.world_size > 1:
+                all_audio_features, all_text_features = gather_features(
+                    audio_features=audio_features,text_features=text_features,
+                    local_loss=self.local_loss,gather_with_grad=self.gather_with_grad,
+                    rank=self.rank,world_size=self.world_size,use_horovod=self.use_horovod,
+                    mlp_loss=self.mlp_loss
+                )
 
+                if self.local_loss:
+                    logits_per_audio = logit_scale_a * audio_features @ all_text_features.T
+                    logits_per_text = logit_scale_a * text_features @ all_audio_features.T
+                else:
+                    logits_per_audio = logit_scale_a * all_audio_features @ all_text_features.T
+                    logits_per_text = logits_per_audio.T
+            else:
+                logits_per_audio = logit_scale_a * audio_features @ text_features.T
+                logits_per_text = logit_scale_a * text_features @ audio_features.T
 
-        total_loss = (
-            F.cross_entropy(a_logits_per_audio, labels) +
-            F.cross_entropy(a_logits_per_text, labels) + 
-            F.cross_entropy(t_logits_per_audio, labels) +
-            F.cross_entropy(t_logits_per_text, labels) 
-            ) / 4
+            # calculated ground-truth and cache if enabled
+            num_logits = logits_per_audio.shape[0]
+            if self.prev_num_logits != num_logits or device not in self.labels:
+                labels = torch.arange(num_logits, device=device, dtype=torch.long)
+                if self.world_size > 1 and self.local_loss:
+                    labels = labels + num_logits * self.rank
+                if self.cache_labels:
+                    self.labels[device] = labels
+                    self.prev_num_logits = num_logits
+            else:
+                labels = self.labels[device]
+            if not self.weighted_loss:
+                total_loss = (
+                    F.cross_entropy(logits_per_audio, labels) +
+                    F.cross_entropy(logits_per_text, labels)
+                    ) / 2
+            else:
+                audio_weight = (all_audio_features@all_audio_features.T).detach()
+                audio_weight = (torch.exp(torch.sum(audio_weight, axis=1)/(self.weight_loss_kappa*len(all_audio_features)))).detach()
+                text_weight = (all_text_features@all_text_features.T).detach()
+                text_weight = (torch.exp(torch.sum(text_weight, axis=1)/(self.weight_loss_kappa*len(all_text_features)))).detach()
+                total_loss = (
+                    F.cross_entropy(logits_per_audio, labels, weight=text_weight) +
+                    F.cross_entropy(logits_per_text, labels, weight=audio_weight)
+                    ) / 2
         return total_loss
+
+def lp_gather_features(
+        pred,
+        target,
+        world_size=1,
+        use_horovod=False
+):
+    if use_horovod:
+        assert hvd is not None, 'Please install horovod'
+        with torch.no_grad():
+            all_preds = hvd.allgather(pred)
+            all_targets = hvd.allgath(target)
+    else:
+        gathered_preds = [torch.zeros_like(pred) for _ in range(world_size)]
+        gathered_targets = [torch.zeros_like(target) for _ in range(world_size)]
+
+        dist.all_gather(gathered_preds, pred)
+        dist.all_gather(gathered_targets, target)
+        all_preds = torch.cat(gathered_preds, dim=0)
+        all_targets = torch.cat(gathered_targets, dim=0)
+
+    return all_preds, all_targets
+
+
+def get_map(pred, target):
+    pred = torch.sigmoid(pred).numpy()
+    target = target.numpy()
+    return np.mean(average_precision_score(target, pred, average=None))
+
+def get_acc(pred, target):
+    pred = torch.argmax(pred,1).numpy()
+    target = torch.argmax(target,1).numpy()
+    return accuracy_score(target, pred)
+
+def get_mauc(pred, target):
+    pred = torch.sigmoid(pred).numpy()
+    target = target.numpy()
+    return np.mean(roc_auc_score(target, pred, average=None))
+
+
+class LPMetrics(object):
+    def __init__(self, metric_names = ['map','acc','mauc']):
+        self.metrics = []
+        for name in metric_names:
+            self.metrics.append(self.get_metric(name))
+        self.metric_names = metric_names
+
+    def get_metric(self,name):
+        if name == 'map':
+            return get_map
+        elif name == 'acc':
+            return get_acc
+        elif name == 'mauc':
+            return get_mauc
+        else:
+            raise ValueError(f'the metric should be at least one of [map, acc, mauc]')
+
+    def evaluate_mertics(self, pred, target):
+        metric_dict = {}
+        for i in range(len(self.metric_names)):
+            metric_dict[self.metric_names[i]] = self.metrics[i](pred, target)
+        return metric_dict
+
+
+def calc_celoss(pred, target):
+    target = torch.argmax(target, 1).long()
+    return nn.CrossEntropyLoss()(pred, target)
+
+
+class LPLoss(nn.Module):
+
+    def __init__(self, loss_name):
+        super().__init__()
+        if loss_name == 'bce':
+            self.loss_func = nn.BCEWithLogitsLoss()
+        elif loss_name == 'ce':
+            self.loss_func = calc_celoss
+        elif loss_name == 'mse':
+            self.loss_func = nn.MSELoss()
+        else:
+            raise ValueError(f'the loss func should be at least one of [bce, ce, mse]')
+
+    def forward(self, pred, target):
+        loss = self.loss_func(pred, target)
+        return loss
+        
\ No newline at end of file
diff --git a/src/open_clip/model.py b/src/open_clip/model.py
index 12d32e8..7af6507 100644
--- a/src/open_clip/model.py
+++ b/src/open_clip/model.py
@@ -20,6 +20,8 @@
 
 from .pann_model import create_pann_model
 from .htsat import create_htsat_model
+from transformers import BertModel, RobertaModel, BartModel
+from transformers.tokenization_utils_base import BatchEncoding
 
 
 class MLPLayers(nn.Module):
@@ -41,6 +43,7 @@ def forward(self, X):
         X = self.sequential(X)
         return X
 
+
 class Bottleneck(nn.Module):
     expansion = 4
 
@@ -65,11 +68,24 @@ def __init__(self, inplanes, planes, stride=1):
 
         if stride > 1 or inplanes != planes * Bottleneck.expansion:
             # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
-            self.downsample = nn.Sequential(OrderedDict([
-                ("-1", nn.AvgPool2d(stride)),
-                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
-                ("1", nn.BatchNorm2d(planes * self.expansion))
-            ]))
+            self.downsample = nn.Sequential(
+                OrderedDict(
+                    [
+                        ("-1", nn.AvgPool2d(stride)),
+                        (
+                            "0",
+                            nn.Conv2d(
+                                inplanes,
+                                planes * self.expansion,
+                                1,
+                                stride=1,
+                                bias=False,
+                            ),
+                        ),
+                        ("1", nn.BatchNorm2d(planes * self.expansion)),
+                    ]
+                )
+            )
 
     def forward(self, x: torch.Tensor):
         identity = x
@@ -88,9 +104,13 @@ def forward(self, x: torch.Tensor):
 
 
 class AttentionPool2d(nn.Module):
-    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+    def __init__(
+        self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None
+    ):
         super().__init__()
-        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+        self.positional_embedding = nn.Parameter(
+            torch.randn(spacial_dim**2 + 1, embed_dim) / embed_dim**0.5
+        )
         self.k_proj = nn.Linear(embed_dim, embed_dim)
         self.q_proj = nn.Linear(embed_dim, embed_dim)
         self.v_proj = nn.Linear(embed_dim, embed_dim)
@@ -98,18 +118,24 @@ def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim:
         self.num_heads = num_heads
 
     def forward(self, x):
-        x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(
+            2, 0, 1
+        )  # NCHW -> (HW)NC
         x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
         x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
         x, _ = F.multi_head_attention_forward(
-            query=x, key=x, value=x,
+            query=x,
+            key=x,
+            value=x,
             embed_dim_to_check=x.shape[-1],
             num_heads=self.num_heads,
             q_proj_weight=self.q_proj.weight,
             k_proj_weight=self.k_proj.weight,
             v_proj_weight=self.v_proj.weight,
             in_proj_weight=None,
-            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            in_proj_bias=torch.cat(
+                [self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]
+            ),
             bias_k=None,
             bias_v=None,
             add_zero_attn=False,
@@ -118,7 +144,7 @@ def forward(self, x):
             out_proj_bias=self.c_proj.bias,
             use_separate_proj_weight=True,
             training=self.training,
-            need_weights=False
+            need_weights=False,
         )
 
         return x[0]
@@ -138,9 +164,13 @@ def __init__(self, layers, output_dim, heads, image_size=224, width=64):
         self.image_size = image_size
 
         # the 3-layer stem
-        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.conv1 = nn.Conv2d(
+            3, width // 2, kernel_size=3, stride=2, padding=1, bias=False
+        )
         self.bn1 = nn.BatchNorm2d(width // 2)
-        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.conv2 = nn.Conv2d(
+            width // 2, width // 2, kernel_size=3, padding=1, bias=False
+        )
         self.bn2 = nn.BatchNorm2d(width // 2)
         self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
         self.bn3 = nn.BatchNorm2d(width)
@@ -170,7 +200,7 @@ def _make_layer(self, planes, blocks, stride=1):
 
     def init_parameters(self):
         if self.attnpool is not None:
-            std = self.attnpool.c_proj.in_features ** -0.5
+            std = self.attnpool.c_proj.in_features**-0.5
             nn.init.normal_(self.attnpool.q_proj.weight, std=std)
             nn.init.normal_(self.attnpool.k_proj.weight, std=std)
             nn.init.normal_(self.attnpool.v_proj.weight, std=std)
@@ -182,14 +212,20 @@ def init_parameters(self):
                     nn.init.zeros_(param)
 
     def lock(self, unlocked_groups=0, freeze_bn_stats=False):
-        assert unlocked_groups == 0, 'partial locking not currently supported for this model'
+        assert (
+            unlocked_groups == 0
+        ), "partial locking not currently supported for this model"
         for param in self.parameters():
             param.requires_grad = False
         if freeze_bn_stats:
             freeze_batch_norm_2d(self)
 
     def stem(self, x):
-        for conv, bn in [(self.conv1, self.bn1), (self.conv2, self.bn2), (self.conv3, self.bn3)]:
+        for conv, bn in [
+            (self.conv1, self.bn1),
+            (self.conv2, self.bn2),
+            (self.conv3, self.bn3),
+        ]:
             x = self.relu(bn(conv(x)))
         x = self.avgpool(x)
         return x
@@ -226,11 +262,15 @@ def __init__(self, d_model: int, n_head: int, act_layer: Callable = nn.GELU):
 
         self.attn = nn.MultiheadAttention(d_model, n_head)
         self.ln_1 = LayerNorm(d_model)
-        self.mlp = nn.Sequential(OrderedDict([
-            ("c_fc", nn.Linear(d_model, d_model * 4)),
-            ("gelu", act_layer()),
-            ("c_proj", nn.Linear(d_model * 4, d_model))
-        ]))
+        self.mlp = nn.Sequential(
+            OrderedDict(
+                [
+                    ("c_fc", nn.Linear(d_model, d_model * 4)),
+                    ("gelu", act_layer()),
+                    ("c_proj", nn.Linear(d_model * 4, d_model)),
+                ]
+            )
+        )
         self.ln_2 = LayerNorm(d_model)
 
     def attention(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
@@ -243,14 +283,18 @@ def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
 
 
 class Transformer(nn.Module):
-    def __init__(self, width: int, layers: int, heads: int, act_layer: Callable = nn.GELU):
+    def __init__(
+        self, width: int, layers: int, heads: int, act_layer: Callable = nn.GELU
+    ):
         super().__init__()
         self.width = width
         self.layers = layers
-        self.resblocks = nn.ModuleList([
-            ResidualAttentionBlock(width, heads, act_layer=act_layer)
-            for _ in range(layers)
-        ])
+        self.resblocks = nn.ModuleList(
+            [
+                ResidualAttentionBlock(width, heads, act_layer=act_layer)
+                for _ in range(layers)
+            ]
+        )
 
     def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
         for r in self.resblocks:
@@ -260,25 +304,42 @@ def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
 
 class VisualTransformer(nn.Module):
     def __init__(
-            self, image_size: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int,
-            act_layer: Callable = nn.GELU):
+        self,
+        image_size: int,
+        patch_size: int,
+        width: int,
+        layers: int,
+        heads: int,
+        output_dim: int,
+        act_layer: Callable = nn.GELU,
+    ):
         super().__init__()
         self.image_size = image_size
         self.output_dim = output_dim
-        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+        self.conv1 = nn.Conv2d(
+            in_channels=3,
+            out_channels=width,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=False,
+        )
 
-        scale = width ** -0.5
+        scale = width**-0.5
         self.class_embedding = nn.Parameter(scale * torch.randn(width))
-        self.positional_embedding = nn.Parameter(scale * torch.randn((image_size // patch_size) ** 2 + 1, width))
+        self.positional_embedding = nn.Parameter(
+            scale * torch.randn((image_size // patch_size) ** 2 + 1, width)
+        )
         self.ln_pre = LayerNorm(width)
 
-        self.transformer = Transformer(width, layers, heads, act_layer=act_layer)
+        self.text_branch = Transformer(width, layers, heads, act_layer=act_layer)
 
         self.ln_post = LayerNorm(width)
         self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
 
     def lock(self, unlocked_groups=0, freeze_bn_stats=False):
-        assert unlocked_groups == 0, 'partial locking not currently supported for this model'
+        assert (
+            unlocked_groups == 0
+        ), "partial locking not currently supported for this model"
         for param in self.parameters():
             param.requires_grad = False
 
@@ -287,13 +348,20 @@ def forward(self, x: torch.Tensor):
         x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
         x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
         x = torch.cat(
-            [self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device),
-             x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+            [
+                self.class_embedding.to(x.dtype)
+                + torch.zeros(
+                    x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device
+                ),
+                x,
+            ],
+            dim=1,
+        )  # shape = [*, grid ** 2 + 1, width]
         x = x + self.positional_embedding.to(x.dtype)
         x = self.ln_pre(x)
 
         x = x.permute(1, 0, 2)  # NLD -> LND
-        x = self.transformer(x)
+        x = self.text_branch(x)
         x = x.permute(1, 0, 2)  # LND -> NLD
 
         x = self.ln_post(x[:, 0, :])
@@ -310,10 +378,19 @@ class CLAPVisionCfg:
     width: int = 768
     patch_size: int = 16
     image_size: Union[Tuple[int, int], int] = 224
-    timm_model_name: str = None  # a valid model name overrides layers, width, patch_size
-    timm_model_pretrained: bool = False  # use (imagenet) pretrained weights for named model
-    timm_pool: str = 'avg'  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
-    timm_proj: str = 'linear'  # linear projection for timm model output ('linear', 'mlp', '')
+    timm_model_name: str = (
+        None  # a valid model name overrides layers, width, patch_size
+    )
+    timm_model_pretrained: bool = (
+        False  # use (imagenet) pretrained weights for named model
+    )
+    timm_pool: str = (
+        "avg"  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
+    )
+    timm_proj: str = (
+        "linear"  # linear projection for timm model output ('linear', 'mlp', '')
+    )
+
 
 # Audio Config Class
 @dataclass
@@ -339,15 +416,20 @@ class CLAPTextCfg:
     width: int
     heads: int
     layers: int
+    model_type: str
 
 
 class CLAP(nn.Module):
     def __init__(
-            self,
-            embed_dim: int,
-            audio_cfg: CLAPAudioCfp,
-            text_cfg: CLAPTextCfg,
-            quick_gelu: bool = False,
+        self,
+        embed_dim: int,
+        audio_cfg: CLAPAudioCfp,
+        text_cfg: CLAPTextCfg,
+        quick_gelu: bool = False,
+        enable_fusion: bool = False,
+        fusion_type: str = 'None',
+        joint_embed_shape: int = 512,
+        mlp_act: str = 'relu',
     ):
         super().__init__()
         if isinstance(audio_cfg, dict):
@@ -357,6 +439,11 @@ def __init__(
 
         self.audio_cfg = audio_cfg
         self.text_cfg = text_cfg
+        self.enable_fusion = enable_fusion
+        self.fusion_type = fusion_type
+        self.joint_embed_shape = joint_embed_shape
+        self.mlp_act = mlp_act
+
 
         self.context_length = text_cfg.context_length
 
@@ -365,70 +452,131 @@ def __init__(
         # NOTE: timm models always use native GELU regardless of quick_gelu flag.
         act_layer = QuickGELU if quick_gelu else nn.GELU
 
+        if mlp_act == 'relu':
+            mlp_act_layer = nn.ReLU()
+        elif mlp_act == 'gelu':
+            mlp_act_layer = nn.GELU()
+        else:
+            raise NotImplementedError
+
         # audio branch
         # audio branch parameters
         if audio_cfg.model_type == "PANN":
-            self.audio_branch = create_pann_model(audio_cfg)
+            self.audio_branch = create_pann_model(audio_cfg, enable_fusion, fusion_type)
         elif audio_cfg.model_type == "HTSAT":
-            self.audio_branch = create_htsat_model(audio_cfg)
+            self.audio_branch = create_htsat_model(audio_cfg, enable_fusion, fusion_type)
         else:
-            logging.error(f'Model config for {audio_cfg.model_type} not found')
-            raise RuntimeError(f'Model config for {audio_cfg.model_type} not found.')
-            
+            logging.error(f"Model config for {audio_cfg.model_type} not found")
+            raise RuntimeError(f"Model config for {audio_cfg.model_type} not found.")
+
         # text branch
         # text branch parameters
-        self.transformer = Transformer(
-            width=text_cfg.width,
-            layers=text_cfg.layers,
-            heads=text_cfg.heads,
-            act_layer=act_layer,
-        )
-
+        if text_cfg.model_type == "transformer":
+            self.text_branch = Transformer(
+                width=text_cfg.width,
+                layers=text_cfg.layers,
+                heads=text_cfg.heads,
+                act_layer=act_layer,
+            )
+            self.vocab_size = text_cfg.vocab_size
+            self.token_embedding = nn.Embedding(text_cfg.vocab_size, text_cfg.width)
+            self.positional_embedding = nn.Parameter(
+                torch.empty(self.context_length, text_cfg.width)
+            )
+            self.ln_final = LayerNorm(text_cfg.width)
+            self.text_transform = MLPLayers(units=[self.joint_embed_shape,
+                                                   self.joint_embed_shape,
+                                                   self.joint_embed_shape], dropout=0.1)
+            self.text_projection = nn.Sequential(
+                nn.Linear(text_cfg.width, self.joint_embed_shape),
+                mlp_act_layer,
+                nn.Linear(self.joint_embed_shape, self.joint_embed_shape)
+            )
+        elif text_cfg.model_type == "bert":
+            self.text_branch = BertModel.from_pretrained("bert-base-uncased")
+            self.text_transform = MLPLayers(units=[self.joint_embed_shape,
+                                                   self.joint_embed_shape,
+                                                   self.joint_embed_shape], dropout=0.1)
+            self.text_projection = nn.Sequential(
+                nn.Linear(768, self.joint_embed_shape),
+                mlp_act_layer,
+                nn.Linear(self.joint_embed_shape, self.joint_embed_shape)
+            )
+        elif text_cfg.model_type == "roberta":
+            self.text_branch = RobertaModel.from_pretrained('roberta-base')
+            self.text_transform = MLPLayers(units=[self.joint_embed_shape,
+                                                   self.joint_embed_shape,
+                                                   self.joint_embed_shape], dropout=0.1)
+            self.text_projection = nn.Sequential(
+                nn.Linear(768, self.joint_embed_shape),
+                mlp_act_layer,
+                nn.Linear(self.joint_embed_shape, self.joint_embed_shape)
+            )
+        elif text_cfg.model_type == "bart":
+            self.text_branch = BartModel.from_pretrained('facebook/bart-base')
+            self.text_transform = MLPLayers(units=[self.joint_embed_shape,
+                                                   self.joint_embed_shape,
+                                                   self.joint_embed_shape], dropout=0.1)
+            self.text_projection = nn.Sequential(
+                nn.Linear(768, self.joint_embed_shape),
+                mlp_act_layer,
+                nn.Linear(self.joint_embed_shape, self.joint_embed_shape)
+            )
+        else:
+            logging.error(f"Model config for {text_cfg.model_type} not found")
+            raise RuntimeError(f"Model config for {text_cfg.model_type} not found.")
+        self.text_branch_type = text_cfg.model_type
         # text branch parameters
-        self.text_transform = MLPLayers(
-            units=[512,512,512], dropout=0.1
-        )
+
         # audio branch parameters
-        self.audio_transform = MLPLayers(
-            units=[512,512,512], dropout=0.1
-        )
+        self.audio_transform = MLPLayers(units=[self.joint_embed_shape,
+                                                self.joint_embed_shape,
+                                                self.joint_embed_shape], dropout=0.1)
 
         # below here is text branch parameters
-        self.vocab_size = text_cfg.vocab_size
-        self.token_embedding = nn.Embedding(text_cfg.vocab_size, text_cfg.width)
-        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, text_cfg.width))
-        self.ln_final = LayerNorm(text_cfg.width)
 
         # ============================================================================================================
-        self.audio_projection = MLPLayers(units=[embed_dim, 512, 512], dropout=0.1)
-        self.text_projection = nn.Parameter(torch.empty(text_cfg.width, 512)) # HARDCORE as 512
+        self.audio_projection = nn.Sequential(
+                nn.Linear(embed_dim, self.joint_embed_shape),
+                mlp_act_layer,
+                nn.Linear(self.joint_embed_shape, self.joint_embed_shape)
+            )
+
         self.logit_scale_a = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
         self.logit_scale_t = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
-        self.register_buffer('attn_mask', self.build_attention_mask(), persistent=False)
+        self.register_buffer("attn_mask", self.build_attention_mask(), persistent=False)
 
         self.init_text_branch_parameters()
 
     def init_text_branch_parameters(self):
-        nn.init.normal_(self.token_embedding.weight, std=0.02)
-        nn.init.normal_(self.positional_embedding, std=0.01)
+        if self.text_branch_type == "transformer":
+            nn.init.normal_(self.token_embedding.weight, std=0.02)
+            nn.init.normal_(self.positional_embedding, std=0.01)
+            proj_std = (self.text_branch.width**-0.5) * (
+                (2 * self.text_branch.layers) ** -0.5
+            )
+            attn_std = self.text_branch.width**-0.5
+            fc_std = (2 * self.text_branch.width) ** -0.5
+            for block in self.text_branch.resblocks:
+                nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+                nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+                nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+                nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+        if self.text_branch_type == "bert" or self.text_branch_type == "roberta":
+            width = self.text_branch.embeddings.word_embeddings.weight.shape[-1]
+        elif self.text_branch_type == "bart":
+            width = self.text_branch.shared.weight.shape[-1]
+        else:
+            width = self.text_branch.width
         nn.init.constant_(self.logit_scale_a, np.log(1 / 0.07))
         nn.init.constant_(self.logit_scale_t, np.log(1 / 0.07))
 
         # deprecated
         # if hasattr(self.visual, 'init_parameters'):
-            # self.visual.init_parameters()
+        # self.visual.init_parameters()
 
-        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
-        attn_std = self.transformer.width ** -0.5
-        fc_std = (2 * self.transformer.width) ** -0.5
-        for block in self.transformer.resblocks:
-            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
-            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
-            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
-            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
-
-        if self.text_projection is not None:
-            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+        # if self.text_projection is not None:
+        #     nn.init.normal_(self.text_projection, std=width**-0.5)
 
     def build_attention_mask(self):
         # lazily create causal attention mask, with full attention between the vision tokens
@@ -438,69 +586,135 @@ def build_attention_mask(self):
         mask.triu_(1)  # zero out the lower diagonal
         return mask
 
-    def encode_audio(self, audio):
-        return self.audio_branch(audio, None) # mix lambda needs to add
-
-    def encode_text(self, text):
-        x = self.token_embedding(text)  # [batch_size, n_ctx, d_model]
-
-        x = x + self.positional_embedding
-        x = x.permute(1, 0, 2)  # NLD -> LND
-        x = self.transformer(x, attn_mask=self.attn_mask)
-        x = x.permute(1, 0, 2)  # LND -> NLD
-        x = self.ln_final(x)
-
-        # x.shape = [batch_size, n_ctx, transformer.width]
-        # take features from the eot embedding (eot_token is the highest number in each sequence)
-        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
-
+    def encode_audio(self, audio, device):
+        return self.audio_branch(audio, mixup_lambda=None, device=device)  # mix lambda needs to add
+
+    # def list_of_dict_of_tensor2dict_of_tensor(self, x, device):
+    #     tmp = {}
+    #     for k in x[0].keys():
+    #         tmp[k] = []
+    #         for i in range(len(x)):
+    #             tmp[k].append(x[i][k][:77])
+    #     for k in x[0].keys():
+    #         tmp[k] = torch.tensor(tmp[k]).to(device=device, non_blocking=True)
+    #     return tmp
+
+    def encode_text(self, text, device):
+        if self.text_branch_type == "transformer":
+            text = text.to(device=device, non_blocking=True)
+            x = self.token_embedding(text)  # [batch_size, n_ctx, d_model]
+
+            x = x + self.positional_embedding
+            x = x.permute(1, 0, 2)  # NLD -> LND
+            x = self.text_branch(x, attn_mask=self.attn_mask)
+            x = x.permute(1, 0, 2)  # LND -> NLD
+            x = self.ln_final(x)
+
+            # x.shape = [batch_size, n_ctx, transformer.width]
+            # take features from the eot embedding (eot_token is the highest number in each sequence)
+            x = self.text_projection(x[torch.arange(x.shape[0]), text.argmax(dim=-1)])
+        elif self.text_branch_type == "bert":
+            # text = self.list_of_dict_of_tensor2dict_of_tensor(text, device)
+            # text = BatchEncoding(text)
+            x = self.text_branch(
+                input_ids=text["input_ids"].to(device=device, non_blocking=True),
+                attention_mask=text["attention_mask"].to(
+                    device=device, non_blocking=True
+                ),
+                token_type_ids=text["token_type_ids"].to(
+                    device=device, non_blocking=True
+                ),
+            )["pooler_output"]
+            x = self.text_projection(x)
+        elif self.text_branch_type == "roberta":
+            x = self.text_branch(
+                input_ids=text["input_ids"].to(device=device, non_blocking=True),
+                attention_mask=text["attention_mask"].to(
+                    device=device, non_blocking=True
+                ),
+            )["pooler_output"]
+            x = self.text_projection(x)
+        elif self.text_branch_type == "bart":
+            x = torch.mean(self.text_branch(
+                input_ids=text["input_ids"].to(device=device, non_blocking=True),
+                attention_mask=text["attention_mask"].to(
+                    device=device, non_blocking=True
+                ),
+            )["encoder_last_hidden_state"],axis=1)
+            x = self.text_projection(x)
+        else:
+            logging.error(f"Model type {self.text_branch_type} not found")
+            raise RuntimeError(f"Model type {self.text_branch_type} not found.")
         return x
 
-    def forward(self, audio, text):
+    def forward(self, audio, text, device=None):
         """Forward audio and text into the CLAP
 
         Parameters
         ----------
         audio: torch.Tensor (batch_size, audio_length)
-            the time-domain audio input
+            the time-domain audio input / the batch of mel_spec and longer list.
         text: torch.Tensor () // need to add
-            the text token input 
+            the text token input
         """
-        if audio is None:
-            return self.encode_text(text)
+        if device is None:
+            if audio is not None:
+                device = audio.device
+            elif text is not None:
+                device = text.device
+        if audio is None and text is None:
+            # a hack to get the logit scale
+            return self.logit_scale_a.exp(), self.logit_scale_t.exp()
+        elif audio is None:
+            return self.encode_text(text, device=device)
         elif text is None:
-            return self.audio_projection(self.encode_audio(audio)["embedding"])
-        audio_features = self.audio_projection(self.encode_audio(audio)["embedding"])
+            return self.audio_projection(self.encode_audio(audio, device=device)["embedding"])
+        audio_features = self.audio_projection(self.encode_audio(audio, device=device)["embedding"])
         audio_features = F.normalize(audio_features, dim=-1)
-        text_features = self.encode_text(text)
+
+        text_features = self.encode_text(
+            text, device=device
+        )
+        # print("text_features", text_features)
+        # print("text_features.shape", text_features.shape)
+        # print("text_features.type", type(text_features))
         text_features = F.normalize(text_features, dim=-1)
 
-        
         audio_features_mlp = self.audio_transform(audio_features)
         text_features_mlp = self.text_transform(text_features)
         # Four outputs: audio features (basic & MLP), text features (basic & MLP)
-        return audio_features, text_features, audio_features_mlp, text_features_mlp, self.logit_scale_a.exp(), self.logit_scale_t.exp()
+        return (
+            audio_features,
+            text_features,
+            audio_features_mlp,
+            text_features_mlp,
+            self.logit_scale_a.exp(),
+            self.logit_scale_t.exp(),
+        )
+
+    def get_logit_scale(self):
+        return self.logit_scale_a.exp(), self.logit_scale_t.exp()
 
-    def audio_infer(self, audio, hopsize = None, key = "embedding"):
+    def audio_infer(self, audio, hopsize=None, key="embedding", device=None):
         """Forward one audio and produce the audio embedding
-        
+
         Parameters
         ----------
         audio: torch.Tensor (audio_length)
             the time-domain audio input, notice that it must be only one input
-        hopsize: int 
+        hopsize: int
             the overlap hopsize as the sliding window
-        key: str 
+        key: str
             the key string to extract the latant code (e.g. embedding, fine_grained_embedding)
 
         Returns
         ----------
         output_dict: {
             key: [n, (embedding_shape)] if "HTS-AT"
-            or 
+            or
             key: [(embedding_shape)] if "PANN"
         }
-            the list of key values of the audio branch 
+            the list of key values of the audio branch
 
         """
 
@@ -508,8 +722,8 @@ def audio_infer(self, audio, hopsize = None, key = "embedding"):
         output_dict = {}
         # PANN
         if self.audio_cfg.model_type == "PANN":
-            audio_input = audio.unsqueeze(dim = 0)
-            output_dict[key] = self.encode_audio(audio_input)[key].squeeze(dim = 0)
+            audio_input = audio.unsqueeze(dim=0)
+            output_dict[key] = self.encode_audio(audio_input, device=device)[key].squeeze(dim=0)
         elif self.audio_cfg.model_type == "HTSAT":
             # repeat
             audio_len = len(audio)
@@ -517,22 +731,25 @@ def audio_infer(self, audio, hopsize = None, key = "embedding"):
             if k > 1:
                 audio = audio.repeat(k)
                 audio_len = len(audio)
-            
+
             if hopsize is None:
                 hopsize = min(hopsize, audio_len)
 
             if audio_len > self.audio_cfg.clip_samples:
-                audio_input = [audio[pos:pos + self.audio_cfg.clip_samples].clone() for pos in range(0, audio_len - self.audio_cfg.clip_samples, hopsize)]
-                audio_input.append(audio[-self.audio_cfg.clip_samples:].clone())
+                audio_input = [
+                    audio[pos : pos + self.audio_cfg.clip_samples].clone()
+                    for pos in range(
+                        0, audio_len - self.audio_cfg.clip_samples, hopsize
+                    )
+                ]
+                audio_input.append(audio[-self.audio_cfg.clip_samples :].clone())
                 audio_input = torch.stack(audio_input)
-                output_dict[key] = self.encode_audio(audio_input)[key]
+                output_dict[key] = self.encode_audio(audio_input, device=device)[key]
             else:
-                audio_input = audio.unsqueeze(dim = 0)
-                output_dict[key] = self.encode_audio(audio_input)[key].squeeze(dim = 0)
-        
-        return output_dict
-            
+                audio_input = audio.unsqueeze(dim=0)
+                output_dict[key] = self.encode_audio(audio_input, device=device)[key].squeeze(dim=0)
 
+        return output_dict
 
 
 def convert_weights_to_fp16(model: nn.Module):
@@ -545,7 +762,12 @@ def _convert_weights_to_fp16(l):
                 l.bias.data = l.bias.data.half()
 
         if isinstance(l, nn.MultiheadAttention):
-            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+            for attr in [
+                *[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]],
+                "in_proj_bias",
+                "bias_k",
+                "bias_v",
+            ]:
                 tensor = getattr(l, attr)
                 if tensor is not None:
                     tensor.data = tensor.data.half()
@@ -559,33 +781,34 @@ def _convert_weights_to_fp16(l):
     model.apply(_convert_weights_to_fp16)
 
 
-
-# Ignore the state dict of the vision part  
-def build_model_from_openai_state_dict(state_dict: dict, model_cfg):
+# Ignore the state dict of the vision part
+def build_model_from_openai_state_dict(state_dict: dict, model_cfg, enable_fusion: bool = False, fusion_type: str = 'None'):
 
     embed_dim = model_cfg["embed_dim"]
     audio_cfg = model_cfg["audio_cfg"]
-
+    text_cfg = model_cfg["text_cfg"]
     context_length = state_dict["positional_embedding"].shape[0]
     vocab_size = state_dict["token_embedding.weight"].shape[0]
     transformer_width = state_dict["ln_final.weight"].shape[0]
     transformer_heads = transformer_width // 64
-    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
-
+    transformer_layers = len(
+        set(
+            k.split(".")[2]
+            for k in state_dict
+            if k.startswith(f"transformer.resblocks")
+        )
+    )
 
     audio_cfg = CLAPAudioCfp(**audio_cfg)
-    text_cfg = CLAPTextCfg(
-        context_length=context_length,
-        vocab_size=vocab_size,
-        width=transformer_width,
-        heads=transformer_heads,
-        layers=transformer_layers
-    )
+    text_cfg = CLAPTextCfg(**text_cfg)
+
     model = CLAP(
         embed_dim,
         audio_cfg=audio_cfg,
         text_cfg=text_cfg,
         quick_gelu=True,  # OpenAI models were trained with QuickGELU
+        enable_fusion=enable_fusion,
+        fusion_type=fusion_type
     )
     state_dict["logit_scale_a"] = state_dict["logit_scale"]
     state_dict["logit_scale_t"] = state_dict["logit_scale"]
@@ -595,7 +818,7 @@ def build_model_from_openai_state_dict(state_dict: dict, model_cfg):
         if key.startswith("visual."):
             state_dict.pop(key, None)
 
-    for key in ["logit_scale","input_resolution", "context_length", "vocab_size"]:
+    for key in ["logit_scale", "input_resolution", "context_length", "vocab_size"]:
         state_dict.pop(key, None)
 
     # not use fp16
@@ -604,17 +827,20 @@ def build_model_from_openai_state_dict(state_dict: dict, model_cfg):
     return model.eval()
 
 
-def trace_model(model, batch_size=256, device=torch.device('cpu')):
+def trace_model(model, batch_size=256, device=torch.device("cpu")):
     model.eval()
     audio_length = model.audio_cfg.audio_length
     example_audio = torch.ones((batch_size, audio_length), device=device)
-    example_text = torch.zeros((batch_size, model.context_length), dtype=torch.int, device=device)
+    example_text = torch.zeros(
+        (batch_size, model.context_length), dtype=torch.int, device=device
+    )
     model = torch.jit.trace_module(
         model,
         inputs=dict(
             forward=(example_audio, example_text),
             encode_text=(example_text,),
-            encode_image=(example_audio,)
-        ))
-    model.audio_cfg.audio_length = audio_length # Question: what does this do?
+            encode_image=(example_audio,),
+        ),
+    )
+    model.audio_cfg.audio_length = audio_length  # Question: what does this do?
     return model
diff --git a/src/open_clip/model_configs/HTSAT-tiny-win-1536.json b/src/open_clip/model_configs/HTSAT-tiny-win-1536.json
new file mode 100644
index 0000000..73e4299
--- /dev/null
+++ b/src/open_clip/model_configs/HTSAT-tiny-win-1536.json
@@ -0,0 +1,23 @@
+{
+    "embed_dim": 768,
+    "audio_cfg": {
+        "audio_length": 1024,
+        "clip_samples": 480000,
+        "mel_bins": 64,
+        "sample_rate": 48000,
+        "window_size": 1536,
+        "hop_size": 480,
+        "fmin": 50,
+        "fmax": 14000,
+        "class_num": 527,
+        "model_type": "HTSAT",
+        "model_name": "tiny"
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}
\ No newline at end of file
diff --git a/src/open_clip/model_configs/PANN-14-fmax-18k.json b/src/open_clip/model_configs/PANN-14-fmax-18k.json
new file mode 100644
index 0000000..b7989bc
--- /dev/null
+++ b/src/open_clip/model_configs/PANN-14-fmax-18k.json
@@ -0,0 +1,23 @@
+{
+    "embed_dim": 2048,
+    "audio_cfg": {
+        "audio_length": 1024,
+        "clip_samples": 480000,
+        "mel_bins": 64,
+        "sample_rate": 48000,
+        "window_size": 1024,
+        "hop_size": 480,
+        "fmin": 50,
+        "fmax": 18000,
+        "class_num": 527,
+        "model_type": "PANN",
+        "model_name": "Cnn14"
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}
\ No newline at end of file
diff --git a/src/open_clip/model_configs/PANN-14-fmax-8k-20s.json b/src/open_clip/model_configs/PANN-14-fmax-8k-20s.json
new file mode 100644
index 0000000..56bdb56
--- /dev/null
+++ b/src/open_clip/model_configs/PANN-14-fmax-8k-20s.json
@@ -0,0 +1,23 @@
+{
+    "embed_dim": 2048,
+    "audio_cfg": {
+        "audio_length": 1024,
+        "clip_samples": 960000,
+        "mel_bins": 64,
+        "sample_rate": 48000,
+        "window_size": 1024,
+        "hop_size": 360,
+        "fmin": 50,
+        "fmax": 8000,
+        "class_num": 527,
+        "model_type": "PANN",
+        "model_name": "Cnn14"
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}
\ No newline at end of file
diff --git a/src/open_clip/model_configs/PANN-14-tiny-transformer.json b/src/open_clip/model_configs/PANN-14-tiny-transformer.json
new file mode 100644
index 0000000..5756e3b
--- /dev/null
+++ b/src/open_clip/model_configs/PANN-14-tiny-transformer.json
@@ -0,0 +1,23 @@
+{
+    "embed_dim": 2048,
+    "audio_cfg": {
+        "audio_length": 1024,
+        "clip_samples": 480000,
+        "mel_bins": 64,
+        "sample_rate": 48000,
+        "window_size": 1024,
+        "hop_size": 480,
+        "fmin": 50,
+        "fmax": 14000,
+        "class_num": 527,
+        "model_type": "PANN",
+        "model_name": "Cnn14"
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 4
+    }
+}
\ No newline at end of file
diff --git a/src/open_clip/model_configs/PANN-14-win-1536.json b/src/open_clip/model_configs/PANN-14-win-1536.json
new file mode 100644
index 0000000..5a9e7e2
--- /dev/null
+++ b/src/open_clip/model_configs/PANN-14-win-1536.json
@@ -0,0 +1,23 @@
+{
+    "embed_dim": 2048,
+    "audio_cfg": {
+        "audio_length": 1024,
+        "clip_samples": 480000,
+        "mel_bins": 64,
+        "sample_rate": 48000,
+        "window_size": 1536,
+        "hop_size": 480,
+        "fmin": 50,
+        "fmax": 14000,
+        "class_num": 527,
+        "model_type": "PANN",
+        "model_name": "Cnn14"
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}
\ No newline at end of file
diff --git a/src/open_clip/openai.py b/src/open_clip/openai.py
index 4ad284d..9911b6e 100644
--- a/src/open_clip/openai.py
+++ b/src/open_clip/openai.py
@@ -26,6 +26,8 @@ def load_openai_model(
         device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu",
         jit=True,
         cache_dir=os.path.expanduser("~/.cache/clip"),
+        enable_fusion: bool = False,
+        fusion_type: str = 'None'
 ):
     """Load a CLIP model, preserve its text pretrained part, and set in the CLAP model
 
@@ -65,10 +67,10 @@ def load_openai_model(
 
     if not jit:
         try:
-            model = build_model_from_openai_state_dict(state_dict or model.state_dict(), model_cfg).to(device)
+            model = build_model_from_openai_state_dict(state_dict or model.state_dict(), model_cfg, enable_fusion, fusion_type).to(device)
         except KeyError:
             sd = {k[7:]: v for k, v in state_dict["state_dict"].items()}
-            model = build_model_from_openai_state_dict(sd, model_cfg).to(device)
+            model = build_model_from_openai_state_dict(sd, model_cfg, enable_fusion, fusion_type).to(device)
 
         if str(device) == "cpu":
             model.float()
diff --git a/src/open_clip/pann_model.py b/src/open_clip/pann_model.py
index ad5efc5..109db5f 100644
--- a/src/open_clip/pann_model.py
+++ b/src/open_clip/pann_model.py
@@ -11,17 +11,18 @@
 from torchlibrosa.augmentation import SpecAugmentation
 
 from .utils import do_mixup, interpolate, pad_framewise_output
- 
+from .feature_fusion import iAFF, AFF, DAF
+
 
 def init_layer(layer):
     """Initialize a Linear or Convolutional layer. """
     nn.init.xavier_uniform_(layer.weight)
- 
+
     if hasattr(layer, 'bias'):
         if layer.bias is not None:
             layer.bias.data.fill_(0.)
-            
-    
+
+
 def init_bn(bn):
     """Initialize a Batchnorm layer. """
     bn.bias.data.fill_(0.)
@@ -30,33 +31,33 @@ def init_bn(bn):
 
 class ConvBlock(nn.Module):
     def __init__(self, in_channels, out_channels):
-        
+
         super(ConvBlock, self).__init__()
-        
-        self.conv1 = nn.Conv2d(in_channels=in_channels, 
+
+        self.conv1 = nn.Conv2d(in_channels=in_channels,
                               out_channels=out_channels,
                               kernel_size=(3, 3), stride=(1, 1),
                               padding=(1, 1), bias=False)
-                              
-        self.conv2 = nn.Conv2d(in_channels=out_channels, 
+
+        self.conv2 = nn.Conv2d(in_channels=out_channels,
                               out_channels=out_channels,
                               kernel_size=(3, 3), stride=(1, 1),
                               padding=(1, 1), bias=False)
-                              
+
         self.bn1 = nn.BatchNorm2d(out_channels)
         self.bn2 = nn.BatchNorm2d(out_channels)
 
         self.init_weight()
-        
+
     def init_weight(self):
         init_layer(self.conv1)
         init_layer(self.conv2)
         init_bn(self.bn1)
         init_bn(self.bn2)
 
-        
+
     def forward(self, input, pool_size=(2, 2), pool_type='avg'):
-        
+
         x = input
         x = F.relu_(self.bn1(self.conv1(x)))
         x = F.relu_(self.bn2(self.conv2(x)))
@@ -70,31 +71,31 @@ def forward(self, input, pool_size=(2, 2), pool_type='avg'):
             x = x1 + x2
         else:
             raise Exception('Incorrect argument!')
-        
+
         return x
 
 
 class ConvBlock5x5(nn.Module):
     def __init__(self, in_channels, out_channels):
-        
+
         super(ConvBlock5x5, self).__init__()
-        
-        self.conv1 = nn.Conv2d(in_channels=in_channels, 
+
+        self.conv1 = nn.Conv2d(in_channels=in_channels,
                               out_channels=out_channels,
                               kernel_size=(5, 5), stride=(1, 1),
                               padding=(2, 2), bias=False)
-                              
+
         self.bn1 = nn.BatchNorm2d(out_channels)
 
         self.init_weight()
-        
+
     def init_weight(self):
         init_layer(self.conv1)
         init_bn(self.bn1)
 
-        
+
     def forward(self, input, pool_size=(2, 2), pool_type='avg'):
-        
+
         x = input
         x = F.relu_(self.bn1(self.conv1(x)))
         if pool_type == 'max':
@@ -107,27 +108,27 @@ def forward(self, input, pool_size=(2, 2), pool_type='avg'):
             x = x1 + x2
         else:
             raise Exception('Incorrect argument!')
-        
+
         return x
 
 
 class AttBlock(nn.Module):
     def __init__(self, n_in, n_out, activation='linear', temperature=1.):
         super(AttBlock, self).__init__()
-        
+
         self.activation = activation
         self.temperature = temperature
         self.att = nn.Conv1d(in_channels=n_in, out_channels=n_out, kernel_size=1, stride=1, padding=0, bias=True)
         self.cla = nn.Conv1d(in_channels=n_in, out_channels=n_out, kernel_size=1, stride=1, padding=0, bias=True)
-        
+
         self.bn_att = nn.BatchNorm1d(n_out)
         self.init_weights()
-        
+
     def init_weights(self):
         init_layer(self.att)
         init_layer(self.cla)
         init_bn(self.bn_att)
-         
+
     def forward(self, x):
         # x: (n_samples, n_in, n_time)
         norm_att = torch.softmax(torch.clamp(self.att(x), -10, 10), dim=-1)
@@ -143,9 +144,9 @@ def nonlinear_transform(self, x):
 
 
 class Cnn14(nn.Module):
-    def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin, 
-        fmax, classes_num):
-        
+    def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
+        fmax, classes_num, enable_fusion=False, fusion_type='None'):
+
         super(Cnn14, self).__init__()
 
         window = 'hann'
@@ -155,23 +156,29 @@ def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
         amin = 1e-10
         top_db = None
 
+        self.enable_fusion = enable_fusion
+        self.fusion_type = fusion_type
+
         # Spectrogram extractor
-        self.spectrogram_extractor = Spectrogram(n_fft=window_size, hop_length=hop_size, 
-            win_length=window_size, window=window, center=center, pad_mode=pad_mode, 
+        self.spectrogram_extractor = Spectrogram(n_fft=window_size, hop_length=hop_size,
+            win_length=window_size, window=window, center=center, pad_mode=pad_mode,
             freeze_parameters=True)
 
         # Logmel feature extractor
-        self.logmel_extractor = LogmelFilterBank(sr=sample_rate, n_fft=window_size, 
-            n_mels=mel_bins, fmin=fmin, fmax=fmax, ref=ref, amin=amin, top_db=top_db, 
+        self.logmel_extractor = LogmelFilterBank(sr=sample_rate, n_fft=window_size,
+            n_mels=mel_bins, fmin=fmin, fmax=fmax, ref=ref, amin=amin, top_db=top_db,
             freeze_parameters=True)
 
         # Spec augmenter
-        self.spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2, 
+        self.spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2,
             freq_drop_width=8, freq_stripes_num=2)
 
         self.bn0 = nn.BatchNorm2d(64)
 
-        self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
+        if (self.enable_fusion) and (self.fusion_type == 'channel_map'):
+            self.conv_block1 = ConvBlock(in_channels=4, out_channels=64)
+        else:
+            self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
         self.conv_block2 = ConvBlock(in_channels=64, out_channels=128)
         self.conv_block3 = ConvBlock(in_channels=128, out_channels=256)
         self.conv_block4 = ConvBlock(in_channels=256, out_channels=512)
@@ -180,32 +187,116 @@ def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
 
         self.fc1 = nn.Linear(2048, 2048, bias=True)
         self.fc_audioset = nn.Linear(2048, classes_num, bias=True)
-        
+
+        if (self.enable_fusion) and (self.fusion_type in ['daf_1d','aff_1d','iaff_1d']):
+            self.mel_conv1d = nn.Sequential(
+                nn.Conv1d(64, 64, kernel_size=5, stride=3, padding=2),
+                nn.BatchNorm1d(64) # No Relu
+            )
+            if self.fusion_type == 'daf_1d':
+                self.fusion_model = DAF()
+            elif self.fusion_type == 'aff_1d':
+                self.fusion_model = AFF(channels=64, type='1D')
+            elif self.fusion_type == 'iaff_1d':
+                self.fusion_model = iAFF(channels=64, type='1D')
+
+        if (self.enable_fusion) and (self.fusion_type in ['daf_2d','aff_2d','iaff_2d']):
+            self.mel_conv2d = nn.Sequential(
+                nn.Conv2d(1, 64, kernel_size=(5,5), stride=(6, 2), padding=(2,2)),
+                nn.BatchNorm2d(64),
+                nn.ReLU(inplace=True)
+            )
+
+            if self.fusion_type == 'daf_2d':
+                self.fusion_model = DAF()
+            elif self.fusion_type == 'aff_2d':
+                self.fusion_model = AFF(channels=64, type='2D')
+            elif self.fusion_type == 'iaff_2d':
+                self.fusion_model = iAFF(channels=64, type='2D')
         self.init_weight()
 
     def init_weight(self):
         init_bn(self.bn0)
         init_layer(self.fc1)
         init_layer(self.fc_audioset)
- 
-    def forward(self, input, mixup_lambda=None):
+
+    def forward(self, input, mixup_lambda=None, device=None):
         """
         Input: (batch_size, data_length)"""
-        x = self.spectrogram_extractor(input)   # (batch_size, 1, time_steps, freq_bins)
-        x = self.logmel_extractor(x)    # (batch_size, 1, time_steps, mel_bins)
 
-        x = x.transpose(1, 3)
-        x = self.bn0(x)
-        x = x.transpose(1, 3)
-        
+        if self.enable_fusion and input["longer"].sum() == 0:
+            # if no audio is longer than 10s, then randomly select one audio to be longer
+            input["longer"][torch.randint(0, input["longer"].shape[0], (1,))] = True
+
+        if not self.enable_fusion:
+            x = self.spectrogram_extractor(input['waveform'].to(device=device, non_blocking=True))   # (batch_size, 1, time_steps, freq_bins)
+            x = self.logmel_extractor(x)    # (batch_size, 1, time_steps, mel_bins)
+
+            x = x.transpose(1, 3)
+            x = self.bn0(x)
+            x = x.transpose(1, 3)
+        else:
+            longer_list = input["longer"].to(device=device, non_blocking=True)
+            x = input["mel_fusion"].to(device=device, non_blocking=True)
+            longer_list_idx = torch.where(longer_list)[0]
+            x = x.transpose(1, 3)
+            x = self.bn0(x)
+            x = x.transpose(1, 3)
+            if self.fusion_type in ['daf_1d','aff_1d','iaff_1d']:
+                new_x = x[:,0:1,:,:].clone().contiguous()
+                # local processing
+                if len(longer_list_idx) > 0:
+                    fusion_x_local = x[longer_list_idx,1:,:,:].clone().contiguous()
+                    FB,FC,FT,FF = fusion_x_local.size()
+                    fusion_x_local = fusion_x_local.view(FB * FC, FT, FF)
+                    fusion_x_local = torch.permute(fusion_x_local, (0,2,1)).contiguous()
+                    fusion_x_local = self.mel_conv1d(fusion_x_local)
+                    fusion_x_local = fusion_x_local.view(FB,FC,FF,fusion_x_local.size(-1))
+                    fusion_x_local = torch.permute(fusion_x_local, (0,2,1,3)).contiguous().flatten(2)
+                    if fusion_x_local.size(-1) < FT:
+                        fusion_x_local = torch.cat([fusion_x_local, torch.zeros((FB,FF,FT- fusion_x_local.size(-1)), device=device)], dim=-1)
+                    else:
+                        fusion_x_local = fusion_x_local[:,:,:FT]
+                    # 1D fusion
+                    new_x = new_x.squeeze(1).permute((0,2,1)).contiguous()
+                    new_x[longer_list_idx] = self.fusion_model(new_x[longer_list_idx], fusion_x_local)
+                    x = new_x.permute((0,2,1)).contiguous()[:,None,:,:]
+                else:
+                    x = new_x
+            elif self.fusion_type in ['daf_2d','aff_2d','iaff_2d','channel_map']:
+                x = x # no change
+
         if self.training:
             x = self.spec_augmenter(x)
-
         # Mixup on spectrogram
         if self.training and mixup_lambda is not None:
             x = do_mixup(x, mixup_lambda)
+        if (self.enable_fusion) and (self.fusion_type in ['daf_2d','aff_2d','iaff_2d']):
+            global_x = x[:,0:1,:,:]
+
+             # global processing
+            B, C, H, W = global_x.shape
+            global_x = self.conv_block1(global_x, pool_size=(2, 2), pool_type='avg')
+            if len(longer_list_idx) > 0:
+                local_x = x[longer_list_idx,1:,:,:].contiguous()
+                TH = global_x.size(-2)
+                # local processing
+                B, C, H, W = local_x.shape
+                local_x = local_x.view(B*C,1,H,W)
+                local_x = self.mel_conv2d(local_x)
+                local_x = local_x.view(B,C,local_x.size(1),local_x.size(2),local_x.size(3))
+                local_x = local_x.permute((0,2,1,3,4)).contiguous().flatten(2,3)
+                TB,TC,_,TW = local_x.size()
+                if local_x.size(-2) < TH:
+                    local_x = torch.cat([local_x, torch.zeros((TB,TC,TH-local_x.size(-2),TW), device=global_x.device)], dim=-2)
+                else:
+                    local_x = local_x[:,:,:TH,:]
+
+                global_x[longer_list_idx] = self.fusion_model(global_x[longer_list_idx],local_x)
+            x = global_x
+        else:
+            x = self.conv_block1(x, pool_size=(2, 2), pool_type='avg')
 
-        x = self.conv_block1(x, pool_size=(2, 2), pool_type='avg')
         x = F.dropout(x, p=0.2, training=self.training)
         x = self.conv_block2(x, pool_size=(2, 2), pool_type='avg')
         x = F.dropout(x, p=0.2, training=self.training)
@@ -218,15 +309,15 @@ def forward(self, input, mixup_lambda=None):
         x = self.conv_block6(x, pool_size=(1, 1), pool_type='avg')
         x = F.dropout(x, p=0.2, training=self.training)
         x = torch.mean(x, dim=3)
-        
+
         latent_x1 = F.max_pool1d(x, kernel_size=3, stride=1, padding=1)
         latent_x2 = F.avg_pool1d(x, kernel_size=3, stride=1, padding=1)
         latent_x = latent_x1 + latent_x2
         latent_x = latent_x.transpose(1, 2)
         latent_x = F.relu_(self.fc1(latent_x))
         latent_output = interpolate(latent_x, 32)
-        
-        
+
+
         (x1, _) = torch.max(x, dim=2)
         x2 = torch.mean(x, dim=2)
         x = x1 + x2
@@ -234,15 +325,15 @@ def forward(self, input, mixup_lambda=None):
         x = F.relu_(self.fc1(x))
         embedding = F.dropout(x, p=0.5, training=self.training)
         clipwise_output = torch.sigmoid(self.fc_audioset(x))
-        
+
         output_dict = {'clipwise_output': clipwise_output, 'embedding': embedding, 'fine_grained_embedding': latent_output}
         return output_dict
 
 
 class Cnn6(nn.Module):
-    def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin, 
-        fmax, classes_num):
-        
+    def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
+        fmax, classes_num, enable_fusion=False, fusion_type='None'):
+
         super(Cnn6, self).__init__()
 
         window = 'hann'
@@ -252,18 +343,21 @@ def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
         amin = 1e-10
         top_db = None
 
+        self.enable_fusion = enable_fusion
+        self.fusion_type = fusion_type
+
         # Spectrogram extractor
-        self.spectrogram_extractor = Spectrogram(n_fft=window_size, hop_length=hop_size, 
-            win_length=window_size, window=window, center=center, pad_mode=pad_mode, 
+        self.spectrogram_extractor = Spectrogram(n_fft=window_size, hop_length=hop_size,
+            win_length=window_size, window=window, center=center, pad_mode=pad_mode,
             freeze_parameters=True)
 
         # Logmel feature extractor
-        self.logmel_extractor = LogmelFilterBank(sr=sample_rate, n_fft=window_size, 
-            n_mels=mel_bins, fmin=fmin, fmax=fmax, ref=ref, amin=amin, top_db=top_db, 
+        self.logmel_extractor = LogmelFilterBank(sr=sample_rate, n_fft=window_size,
+            n_mels=mel_bins, fmin=fmin, fmax=fmax, ref=ref, amin=amin, top_db=top_db,
             freeze_parameters=True)
 
         # Spec augmenter
-        self.spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2, 
+        self.spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2,
             freq_drop_width=8, freq_stripes_num=2)
 
         self.bn0 = nn.BatchNorm2d(64)
@@ -275,32 +369,32 @@ def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
 
         self.fc1 = nn.Linear(512, 512, bias=True)
         self.fc_audioset = nn.Linear(512, classes_num, bias=True)
-        
+
         self.init_weight()
 
     def init_weight(self):
         init_bn(self.bn0)
         init_layer(self.fc1)
         init_layer(self.fc_audioset)
- 
-    def forward(self, input, mixup_lambda=None):
+
+    def forward(self, input, mixup_lambda=None, device=None):
         """
         Input: (batch_size, data_length)"""
 
         x = self.spectrogram_extractor(input)   # (batch_size, 1, time_steps, freq_bins)
         x = self.logmel_extractor(x)    # (batch_size, 1, time_steps, mel_bins)
-        
+
         x = x.transpose(1, 3)
         x = self.bn0(x)
         x = x.transpose(1, 3)
-        
+
         if self.training:
             x = self.spec_augmenter(x)
 
         # Mixup on spectrogram
         if self.training and mixup_lambda is not None:
             x = do_mixup(x, mixup_lambda)
-        
+
         x = self.conv_block1(x, pool_size=(2, 2), pool_type='avg')
         x = F.dropout(x, p=0.2, training=self.training)
         x = self.conv_block2(x, pool_size=(2, 2), pool_type='avg')
@@ -310,14 +404,14 @@ def forward(self, input, mixup_lambda=None):
         x = self.conv_block4(x, pool_size=(2, 2), pool_type='avg')
         x = F.dropout(x, p=0.2, training=self.training)
         x = torch.mean(x, dim=3)
-        
+
         latent_x1 = F.max_pool1d(x, kernel_size=3, stride=1, padding=1)
         latent_x2 = F.avg_pool1d(x, kernel_size=3, stride=1, padding=1)
         latent_x = latent_x1 + latent_x2
         latent_x = latent_x.transpose(1, 2)
         latent_x = F.relu_(self.fc1(latent_x))
         latent_output = interpolate(latent_x, 16)
-        
+
         (x1, _) = torch.max(x, dim=2)
         x2 = torch.mean(x, dim=2)
         x = x1 + x2
@@ -325,16 +419,16 @@ def forward(self, input, mixup_lambda=None):
         x = F.relu_(self.fc1(x))
         embedding = F.dropout(x, p=0.5, training=self.training)
         clipwise_output = torch.sigmoid(self.fc_audioset(x))
-        
+
         output_dict = {'clipwise_output': clipwise_output, 'embedding': embedding, 'fine_grained_embedding': latent_output}
 
         return output_dict
 
 
 class Cnn10(nn.Module):
-    def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin, 
-        fmax, classes_num):
-        
+    def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
+        fmax, classes_num, enable_fusion=False, fusion_type='None'):
+
         super(Cnn10, self).__init__()
 
         window = 'hann'
@@ -344,18 +438,21 @@ def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
         amin = 1e-10
         top_db = None
 
+        self.enable_fusion = enable_fusion
+        self.fusion_type = fusion_type
+
         # Spectrogram extractor
-        self.spectrogram_extractor = Spectrogram(n_fft=window_size, hop_length=hop_size, 
-            win_length=window_size, window=window, center=center, pad_mode=pad_mode, 
+        self.spectrogram_extractor = Spectrogram(n_fft=window_size, hop_length=hop_size,
+            win_length=window_size, window=window, center=center, pad_mode=pad_mode,
             freeze_parameters=True)
 
         # Logmel feature extractor
-        self.logmel_extractor = LogmelFilterBank(sr=sample_rate, n_fft=window_size, 
-            n_mels=mel_bins, fmin=fmin, fmax=fmax, ref=ref, amin=amin, top_db=top_db, 
+        self.logmel_extractor = LogmelFilterBank(sr=sample_rate, n_fft=window_size,
+            n_mels=mel_bins, fmin=fmin, fmax=fmax, ref=ref, amin=amin, top_db=top_db,
             freeze_parameters=True)
 
         # Spec augmenter
-        self.spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2, 
+        self.spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2,
             freq_drop_width=8, freq_stripes_num=2)
 
         self.bn0 = nn.BatchNorm2d(64)
@@ -368,32 +465,32 @@ def __init__(self, sample_rate, window_size, hop_size, mel_bins, fmin,
 
         self.fc1 = nn.Linear(1024, 1024, bias=True)
         self.fc_audioset = nn.Linear(1024, classes_num, bias=True)
-        
+
         self.init_weight()
 
     def init_weight(self):
         init_bn(self.bn0)
         init_layer(self.fc1)
         init_layer(self.fc_audioset)
- 
-    def forward(self, input, mixup_lambda=None):
+
+    def forward(self, input, mixup_lambda=None, device=None):
         """
         Input: (batch_size, data_length)"""
 
         x = self.spectrogram_extractor(input)   # (batch_size, 1, time_steps, freq_bins)
         x = self.logmel_extractor(x)    # (batch_size, 1, time_steps, mel_bins)
-        
+
         x = x.transpose(1, 3)
         x = self.bn0(x)
         x = x.transpose(1, 3)
-        
+
         if self.training:
             x = self.spec_augmenter(x)
 
         # Mixup on spectrogram
         if self.training and mixup_lambda is not None:
             x = do_mixup(x, mixup_lambda)
-        
+
         x = self.conv_block1(x, pool_size=(2, 2), pool_type='avg')
         x = F.dropout(x, p=0.2, training=self.training)
         x = self.conv_block2(x, pool_size=(2, 2), pool_type='avg')
@@ -405,14 +502,14 @@ def forward(self, input, mixup_lambda=None):
         x = self.conv_block5(x, pool_size=(2, 2), pool_type='avg')
         x = F.dropout(x, p=0.2, training=self.training)
         x = torch.mean(x, dim=3)
-        
+
         latent_x1 = F.max_pool1d(x, kernel_size=3, stride=1, padding=1)
         latent_x2 = F.avg_pool1d(x, kernel_size=3, stride=1, padding=1)
         latent_x = latent_x1 + latent_x2
         latent_x = latent_x.transpose(1, 2)
         latent_x = F.relu_(self.fc1(latent_x))
         latent_output = interpolate(latent_x, 32)
-        
+
         (x1, _) = torch.max(x, dim=2)
         x2 = torch.mean(x, dim=2)
         x = x1 + x2
@@ -420,13 +517,13 @@ def forward(self, input, mixup_lambda=None):
         x = F.relu_(self.fc1(x))
         embedding = F.dropout(x, p=0.5, training=self.training)
         clipwise_output = torch.sigmoid(self.fc_audioset(x))
-        
+
         output_dict = {'clipwise_output': clipwise_output, 'embedding': embedding, 'fine_grained_embedding': latent_output}
 
         return output_dict
 
 
-def create_pann_model(audio_cfg):
+def create_pann_model(audio_cfg, enable_fusion=False, fusion_type='None'):
     try:
         ModelProto = eval(audio_cfg.model_name)
         model = ModelProto(
@@ -436,9 +533,11 @@ def create_pann_model(audio_cfg):
             mel_bins = audio_cfg.mel_bins,
             fmin = audio_cfg.fmin,
             fmax = audio_cfg.fmax,
-            classes_num = audio_cfg.class_num
+            classes_num = audio_cfg.class_num,
+            enable_fusion = enable_fusion,
+            fusion_type = fusion_type
         )
         return model
     except:
         raise RuntimeError(f'Import Model for {audio_cfg.model_name} not found, or the audio cfg parameters are not enough.')
-        
+
diff --git a/src/open_clip/utils.py b/src/open_clip/utils.py
index 0769913..ad9801c 100644
--- a/src/open_clip/utils.py
+++ b/src/open_clip/utils.py
@@ -8,8 +8,9 @@
 import random
 import json
 import os
+import pathlib
 
-
+# TODO: (yusong) this not a good place to store those information and does not scale. Need to be fixed later.
 dataset_split = {
     "audiocaps": ["train", "valid", "test"],
     "audioset": ["balanced_train", "unbalanced_train", "eval"],
@@ -19,6 +20,22 @@
     "paramount_motion": ["train", "test"],
     "sonniss_game_effects": ["train", "test"],
     "wesoundeffects": ["train", "test"],
+    "MACS": ["train", "test"],
+    "freesound": ["train", "test"],
+    "FSD50K": ["train", "test", "valid"],
+    "fsd50k_class_label": ["train", "test", "valid"],
+    "esc50": ["train", "test"],
+    "audiostock": ["train", "test"],
+    "freesound_no_overlap_noesc50": ["train", "test"],
+    "epidemic_sound_effects": ["train", "test"],
+    "VGGSound": ["train", "test"],
+    "urbansound8k_class_label": ["train", "test"],
+    "audioset_t5": ["balanced_train", "unbalanced_train", "eval"],
+    "epidemic_sound_effects_t5": ["train", "test"],
+    "WavText5K": ["train", "test"],
+    "esc50_no_overlap": ["train", "test"],
+    "usd8k_no_overlap": ["train", "test"],
+    "fsd50k_200_class_label": ["train", "test", "valid"]
 }
 
 
@@ -43,7 +60,7 @@ def freeze_batch_norm_2d(module, module_match={}, name=""):
     if module_match:
         is_match = name in module_match
     if is_match and isinstance(
-        module, (nn.modules.batchnorm.BatchNorm2d, nn.modules.batchnorm.SyncBatchNorm)
+            module, (nn.modules.batchnorm.BatchNorm2d, nn.modules.batchnorm.SyncBatchNorm)
     ):
         res = FrozenBatchNorm2d(module.num_features)
         res.num_features = module.num_features
@@ -74,68 +91,44 @@ def exist(dataset_name, dataset_type):
 
 
 def get_tar_path_from_dataset_name(
-    dataset_names,
-    dataset_types,
-    islocal,
-    dataset_path,
-    proportion=1,
+        dataset_names,
+        dataset_types,
+        islocal,
+        dataset_path,
+        proportion=1,
+        full_dataset=None
 ):
     """
     Get tar path from dataset name and type
     """
-    if islocal:
-        # txt_paths = []
-        # for i in range(len(dataset_names)):
-        #     for j in range(len(dataset_types)):
-        #         if exist(dataset_names[i], dataset_types[j]):
-        #             txt_loc = template.replace("datasetname", dataset_names[i]).replace(
-        #                 "datasettype", dataset_types[j]
-        #             )
-        #             txt_paths.append(txt_loc)
-        #         else:
-        #             print(
-        #                 "Skipping dataset "
-        #                 + dataset_names[i]
-        #                 + " with "
-        #                 + dataset_types[j]
-        #             )
-        #             continue
-        # return get_tar_path_from_txts(txt_paths, islocal=islocal, proportion=proportion)
-        output = []
-        for n in dataset_names:
-            for s in dataset_types:
-                tmp = []
-                sizefilepath_ = f"./json_files/{n}/{s}/sizes.json" #  TODO:!!!
+    output = []
+    for n in dataset_names:
+        if full_dataset is not None and n in full_dataset:
+            current_dataset_types = dataset_split[n]
+        else:
+            current_dataset_types = dataset_types
+        for s in current_dataset_types:
+            tmp = []
+            if islocal:
+                sizefilepath_ = f"{dataset_path}/{n}/{s}/sizes.json"
                 if not os.path.exists(sizefilepath_):
-                    continue
-                sizes = json.load(open(sizefilepath_, "r"))
-                for k in sizes.keys():
-                    tmp.append(
-                        f"{dataset_path}/{n}/{s}/{k}"
-                    )
-                if proportion!=1:
-                    tmp = random.sample(tmp, int(proportion * len(tmp)))
-                output.append(tmp)
-        return sum(output, [])
-    else:
-
-        output = []
-        for n in dataset_names:
-            for s in dataset_types:
-                tmp = []
+                    sizefilepath_ = f"./json_files/{n}/{s}/sizes.json"
+            else:
                 sizefilepath_ = f"./json_files/{n}/{s}/sizes.json"
-                if not os.path.exists(sizefilepath_):
-                    continue
-                sizes = json.load(open(sizefilepath_, "r"))
-                for k in sizes.keys():
+            if not os.path.exists(sizefilepath_):
+                continue
+            sizes = json.load(open(sizefilepath_, "r"))
+            for k in sizes.keys():
+                if islocal:
+                    tmp.append(f"{dataset_path}/{n}/{s}/{k}")
+                else:
                     tmp.append(
                         f"pipe:aws s3 --cli-connect-timeout 0 cp s3://s-laion-audio/webdataset_tar/{n}/{s}/{k} -"
                     )
-                    # TODO: add dataset_path to remote dataset in the future.
-                if proportion!=1:
-                    tmp = random.sample(tmp, int(proportion * len(tmp)))
-                output.append(tmp)
-        return sum(output, [])
+            if proportion != 1:
+                tmp = random.sample(tmp, int(proportion * len(tmp)))
+            output.append(tmp)
+    return sum(output, [])
 
 
 def get_tar_path_from_txts(txt_path, islocal, proportion=1):
@@ -189,8 +182,8 @@ def do_mixup(x, mixup_lambda):
       out: (batch_size, ...)
     """
     out = (
-        x.transpose(0, -1) * mixup_lambda
-        + torch.flip(x, dims=[0]).transpose(0, -1) * (1 - mixup_lambda)
+            x.transpose(0, -1) * mixup_lambda
+            + torch.flip(x, dims=[0]).transpose(0, -1) * (1 - mixup_lambda)
     ).transpose(0, -1)
     return out
 
@@ -285,3 +278,92 @@ def get_data_from_log(txt_path):
             "train_loss": train_losses[i],
         }
     return train_data, val_data
+
+
+def save_p(obj, filename):
+    import pickle
+
+    try:
+        from deepdiff import DeepDiff
+    except:
+        os.system("pip install deepdiff")
+        from deepdiff import DeepDiff
+    with open(filename, "wb") as file:
+        pickle.dump(obj, file, protocol=pickle.HIGHEST_PROTOCOL)  # highest protocol
+    with open(filename, "rb") as file:
+        z = pickle.load(file)
+    assert (
+            DeepDiff(obj, z, ignore_string_case=True) == {}
+    ), "there is something wrong with the saving process"
+    return
+
+
+def load_p(filename):
+    import pickle
+
+    with open(filename, "rb") as file:
+        z = pickle.load(file)
+    return z
+
+
+def save_json(data, name="data.json"):
+    import json
+    with open(name, 'w') as fp:
+        json.dump(data, fp)
+    return
+
+
+def load_json(name):
+    import json
+    with open(name, 'r') as fp:
+        data = json.load(fp)
+    return data
+
+
+from multiprocessing import Process, Manager
+from multiprocessing import Process, Value, Array
+from ctypes import c_wchar
+
+
+def load_class_label(path):
+    # https://stackoverflow.com/questions/48004243/how-to-share-large-read-only-dictionary-list-across-processes-in-multiprocessing
+    # https://stackoverflow.com/questions/45693949/storing-strings-in-a-multiprocessing-sharedctypes-array
+    out = None
+    if path is not None:
+        if pathlib.Path(path).suffix in [".pkl", ".pickle"]:
+            out = load_p(path)
+        elif pathlib.Path(path).suffix in [".json", ".txt"]:
+            out = load_json(path)
+        elif pathlib.Path(path).suffix in [".npy", ".npz"]:
+            out = np.load(path)
+        elif pathlib.Path(path).suffix in [".csv"]:
+            import pandas as pd
+            out = pd.read_csv(path)
+    return out
+    # if out is None:
+    #     return None
+    # else:
+    #     key = Array(c_wchar, '\n'.join(list(out.keys())), lock=False)
+    #     val = Array('i', out.values(), lock=False)
+    #     return (key, val)
+
+
+from torch import optim
+
+
+def get_optimizer(params, lr, betas, eps, momentum, optimizer_name):
+    if optimizer_name.lower() == "adamw":
+        optimizer = optim.AdamW(
+            params, lr=lr, betas=betas, eps=eps
+        )
+    elif optimizer_name.lower() == "sgd":
+        optimizer = optim.SGD(
+            params, lr=lr, momentum=momentum
+        )
+    elif optimizer_name.lower() == "adam":
+        optimizer = optim.Adam(
+            params, lr=lr, betas=betas, eps=eps
+        )
+    else:
+        raise ValueError("optimizer name is not correct")
+    return optimizer
diff --git a/src/training/data.py b/src/training/data.py
index 037d00c..a833c4d 100644
--- a/src/training/data.py
+++ b/src/training/data.py
@@ -6,12 +6,15 @@
 import random
 import h5py
 from dataclasses import dataclass
-
+from training.params import parse_args
 import braceexpand
 import numpy as np
 import pandas as pd
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
 import torchvision.datasets as datasets
+import torchvision.transforms
 import webdataset as wds
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler
@@ -21,16 +24,70 @@
 import io
 from pathlib import Path
 import wget
-from open_clip import tokenize
-from open_clip.utils import dataset_split
+
+from open_clip.utils import get_tar_path_from_dataset_name, dataset_split
+from open_clip.utils import load_p, load_class_label
+import tempfile
+import copy
 
 try:
     import horovod.torch as hvd
 except ImportError:
     hvd = None
 
-
-from open_clip import tokenize
+try:
+    import torchaudio
+except ImportError:
+    torchaudio = None
+
+args = parse_args()
+if args.tmodel == "transformer":
+    from open_clip import tokenize
+
+    def tokenizer(text):
+        return tokenize(text).squeeze(0)
+
+elif args.tmodel == "bert":
+    from transformers import BertTokenizer
+    tokenize = BertTokenizer.from_pretrained("bert-base-uncased")
+
+    def tokenizer(text):
+        result = tokenize(
+            text,
+            padding="max_length",
+            truncation=True,
+            max_length=77,
+            return_tensors="pt",
+        )
+        return {k: v.squeeze(0) for k, v in result.items()}
+
+elif args.tmodel == "roberta":
+    from transformers import RobertaTokenizer
+    tokenize = RobertaTokenizer.from_pretrained('roberta-base')
+
+    def tokenizer(text):
+        result = tokenize(
+            text,
+            padding="max_length",
+            truncation=True,
+            max_length=77,
+            return_tensors="pt",
+        )
+        return {k: v.squeeze(0) for k, v in result.items()}
+
+elif args.tmodel == "bart":
+    from transformers import BartTokenizer
+    tokenize = BartTokenizer.from_pretrained('facebook/bart-base')
+
+    def tokenizer(text):
+        result = tokenize(
+            text,
+            padding="max_length",
+            truncation=True,
+            max_length=77,
+            return_tensors="pt",
+        )
+        return {k: v.squeeze(0) for k, v in result.items()}
 
 # initizlied the audioset map
 _AUDIOSET_MAP_PATH = os.path.join(Path(__file__).parent, "audioset_textmap.npy")
@@ -41,11 +98,15 @@ def int16_to_float32(x):
     return (x / 32767.0).astype(np.float32)
 
 
+def float32_to_int16(x):
+    x = np.clip(x, a_min=-1., a_max=1.)
+    return (x * 32767.).astype(np.int16)
+
+
 # For Toy Dataset
 class ToyDataset(Dataset):
     def __init__(self, index_path, ipc, config, eval_mode=False):
         """Toy Dataset for testing the audioset input with text labels
-
         Parameters
         ----------
             index_path: str
@@ -132,8 +193,8 @@ def __getitem__(self, index):
             self.fp["hdf5_path"][s_index]
             .decode()
             .replace(
-                "/home/tiger/DB/knut/data/audioset/hdf5s/waveforms",
-                "/mnt/audio_clip/test/data",
+                "../workspace",
+                "/home/la/kechen/Research/ke_zsasp/workspace",
             )
         )
         r_idx = self.fp["index_in_hdf5"][s_index]
@@ -163,14 +224,20 @@ def __getitem__(self, index):
         #             target[add_key] = 1.0
 
         # missing the text input
-
+        mel_spec = get_mel(torch.from_numpy(waveform), self.audio_cfg)[None,:,:]
+        mel_spec = torch.cat([mel_spec, mel_spec.clone(),mel_spec.clone(),mel_spec.clone()], dim=0).cpu().numpy()
+        longer = random.choice([True, False])
+        if longer == False:
+            mel_spec[1:,:,:] = 0.0
         data_dict = {
             "hdf5_path": hdf5_path,
             "index_in_hdf5": r_idx,
             "audio_name": audio_name,
             "waveform": waveform,
-            "target": target,
+            "class_label": target,
             "text": text,
+            "longer": longer,
+            "mel_fusion": mel_spec
         }
         return data_dict
 
@@ -362,51 +429,210 @@ def sample_prop(sizefile, inputs, proportion, is_local=True):
         sampled_size_dict,
     )
 
+def get_mel(audio_data,audio_cfg):
+    # mel shape: (n_mels, T)
+    mel = torchaudio.transforms.MelSpectrogram(
+        sample_rate=audio_cfg['sample_rate'],
+        n_fft=audio_cfg['window_size'],
+        win_length=audio_cfg['window_size'],
+        hop_length=audio_cfg['hop_size'],
+        center=True,
+        pad_mode="reflect",
+        power=2.0,
+        norm=None,
+        onesided=True,
+        n_mels=64,
+        f_min=audio_cfg['fmin'],
+        f_max=audio_cfg['fmax']
+    )(audio_data)
+    # Align to librosa:
+    # librosa_melspec = librosa.feature.melspectrogram(
+    #     waveform,
+    #     sr=audio_cfg['sample_rate'],
+    #     n_fft=audio_cfg['window_size'],
+    #     hop_length=audio_cfg['hop_size'],
+    #     win_length=audio_cfg['window_size'],
+    #     center=True,
+    #     pad_mode="reflect",
+    #     power=2.0,
+    #     n_mels=64,
+    #     norm=None,
+    #     htk=True,
+    #     f_min=audio_cfg['fmin'],
+    #     f_max=audio_cfg['fmax']
+    # )
+    # we use log mel spectrogram as input
+    mel = torchaudio.transforms.AmplitudeToDB(top_db=None)(mel)
+    return mel.T  # (T, n_mels)
+
 
 def preprocess(
     sample,
     audio_ext,
     text_ext,
-    samplerate,
-    mono,
     max_len,
-    dtype,
-    res_type,
-    resample_method="TorchAudio",
+    audio_cfg,
+    class_index_dict=None,
+    data_filling="pad",
+    data_truncating="rand_trunc",
+    text_augment_selection=None,
 ):
     """
     Preprocess a single sample for wdsdataloader.
     """
     audio_data, orig_sr = sf.read(io.BytesIO(sample[audio_ext]))
+    audio_data = int16_to_float32(float32_to_int16(audio_data))
+    audio_data = torch.tensor(audio_data).float()
+
+    # TODO: (yusong) to be include in the future
+    # # if torchaudio not installed, use soundfile to load audio
+    # if torchaudio is None:
+    #     audio_data, orig_sr = sf.read(io.BytesIO(sample[audio_ext]))
+    #     audio_data = torch.tensor(audio_data).float()
+    # else:
+    #     # https://github.com/webdataset/webdataset/blob/main/webdataset/autodecode.py
+    #     with tempfile.TemporaryDirectory() as dirname:
+    #         os.makedirs(dirname, exist_ok=True)
+    #         fname = os.path.join(dirname, f"file.flac")
+    #         with open(fname, "wb") as stream:
+    #             stream.write(sample[audio_ext])
+    #         audio_data, orig_sr = torchaudio.load(fname)
+    #         audio_data = audio_data[0, :].float()
+
+    with torch.no_grad():
+        if len(audio_data) > max_len:
+            if data_truncating == "rand_trunc":
+                longer = torch.tensor([True])
+            elif data_truncating == "fusion":
+                # fusion
+                mel = get_mel(audio_data, audio_cfg)
+                # split to three parts
+                chunk_frames = max_len // audio_cfg['hop_size']+1  # the +1 related to how the spectrogram is computed
+                total_frames = mel.shape[0]
+                if chunk_frames == total_frames:
+                    # there is a corner case where the audio length is
+                    # larger than max_len but smaller than max_len+hop_size.
+                    # In this case, we just use the whole audio.
+                    mel_fusion = torch.stack([mel, mel, mel, mel], dim=0)
+                    sample["mel_fusion"] = mel_fusion
+                    longer = torch.tensor([False])
+                else:
+                    ranges = np.array_split(list(range(0, total_frames-chunk_frames+1)), 3)
+                    # print('total_frames-chunk_frames:', total_frames-chunk_frames,
+                    #       'len(audio_data):', len(audio_data),
+                    #       'chunk_frames:', chunk_frames,
+                    #       'total_frames:', total_frames)
+                    if len(ranges[1]) == 0:
+                        # if the audio is too short, we just use the first chunk
+                        ranges[1] = [0]
+                    if len(ranges[2]) == 0:
+                        # if the audio is too short, we just use the first chunk
+                        ranges[2] = [0]
+                    # randomly choose index for each part
+                    idx_front = np.random.choice(ranges[0])
+                    idx_middle = np.random.choice(ranges[1])
+                    idx_back = np.random.choice(ranges[2])
+                    # select mel
+                    mel_chunk_front = mel[idx_front:idx_front+chunk_frames, :]
+                    mel_chunk_middle = mel[idx_middle:idx_middle+chunk_frames, :]
+                    mel_chunk_back = mel[idx_back:idx_back+chunk_frames, :]
+
+                    # shrink the mel
+                    mel_shrink = torchvision.transforms.Resize(size=[chunk_frames, 64])(mel[None])[0]
+                    # logging.info(f"mel_shrink.shape: {mel_shrink.shape}")
+
+                    # stack
+                    mel_fusion = torch.stack([mel_chunk_front, mel_chunk_middle, mel_chunk_back, mel_shrink], dim=0)
+                    sample["mel_fusion"] = mel_fusion
+                    longer = torch.tensor([True])
+            else:
+                raise NotImplementedError(
+                    f"data_truncating {data_truncating} not implemented"
+                )
+            # random crop to max_len (for compatibility)
+            overflow = len(audio_data) - max_len
+            idx = np.random.randint(0, overflow + 1)
+            audio_data = audio_data[idx: idx + max_len]
+
+        else:  # padding if too short
+            if len(audio_data) < max_len:  # do nothing if equal
+                if data_filling == "repeatpad":
+                    n_repeat = int(max_len/len(audio_data))
+                    audio_data = audio_data.repeat(n_repeat)
+                    # audio_data = audio_data.unsqueeze(0).unsqueeze(0).unsqueeze(0)
+                    # audio_data = F.interpolate(audio_data,size=max_len,mode="bicubic")[0,0,0]
+                    audio_data = F.pad(
+                        audio_data,
+                        (0, max_len - len(audio_data)),
+                        mode="constant",
+                        value=0,
+                    )
+                elif data_filling == "pad":
+                    audio_data = F.pad(
+                        audio_data,
+                        (0, max_len - len(audio_data)),
+                        mode="constant",
+                        value=0,
+                    )
+                elif data_filling == "repeat":
+                    n_repeat = int(max_len/len(audio_data))
+                    audio_data = audio_data.repeat(n_repeat+1)[:max_len]
+                else:
+                    raise NotImplementedError(
+                        f"data_filling {data_filling} not implemented"
+                    )
+            if data_truncating == 'fusion':
+                mel = get_mel(audio_data, audio_cfg)
+                mel_fusion = torch.stack([mel, mel, mel, mel], dim=0)
+                sample["mel_fusion"] = mel_fusion
+            longer = torch.tensor([False])
+
+    sample["longer"] = longer
+
+    sample["waveform"] = audio_data
+    del sample[audio_ext]
 
-    if len(audio_data) > max_len:  # random clip if too long
-        overflow = len(audio_data) - max_len
-        idx = np.random.randint(0, overflow + 1)
-        audio_data = audio_data[idx : idx + max_len]
-        # if np.random.rand() > 0.5:
-        #    audio_data = audio_data[idx : idx + max_len]
-        # else:
-        #    audio_data = audio_data[
-        #        len(audio_data) + 1 - idx - max_len : len(audio_data) + 1 - idx
-        #    ]
-    else:  # padding if too short
-        audio_data = np.pad(
-            audio_data,
-            (0, max_len - len(audio_data)),
-            mode="constant",
-            constant_values=0,
+    try:
+        json_dict_raw = json.loads(sample[text_ext].decode("utf-8"))
+    except:
+        print("sample[__url__]:", sample["__url__"])
+
+    # For selecting augmented text from dataset
+    if text_augment_selection is None or text_augment_selection == "none":
+        texts = json_dict_raw["text"]
+    elif text_augment_selection == "all":
+        if "text_augment_all" in json_dict_raw.keys():
+            texts = json_dict_raw["text_augment_all"]
+        else:
+            texts = json_dict_raw["text"]
+    elif text_augment_selection == "augment_only":
+        if "text_augment_all" in json_dict_raw.keys():
+            if json_dict_raw["text_augment_t5"] is None:
+                texts = json_dict_raw["text"]
+            else:
+                texts = json_dict_raw["text_augment_t5"]
+        else:
+            texts = json_dict_raw["text"]
+    else:
+        raise NotImplementedError(
+            f"text_augment_selection {text_augment_selection} not implemented"
         )
+    sample["full_text"] = texts
 
-    # TODO: (yusong) add a key of "original audio"
-    # TODO: (yusong) also return "original data"
-
-    sample["waveform"] = torch.tensor(audio_data).float()
-    del sample[audio_ext]
-    texts = json.loads(sample[text_ext].decode("utf-8"))["text"]
     if isinstance(texts, list) and isinstance(texts[0], str) and len(texts) > 1:
         texts = random.choice(texts)
     sample["raw_text"] = texts
-    sample["text"] = tokenize(texts)
+    sample["text"] = tokenizer(texts)  # text shape: [num_token]
+    if class_index_dict is not None:
+        # https://stackoverflow.com/questions/48004243/how-to-share-large-read-only-dictionary-list-across-processes-in-multiprocessing
+        # https://stackoverflow.com/questions/45693949/storing-strings-in-a-multiprocessing-sharedctypes-array
+        # key, val = class_index_dict
+        # key = key[:].split('\n')
+        # _dict = {k: v for k, v in zip(key, val)}
+        sample["class_label"] = np.zeros(len(class_index_dict.keys()))
+        for x in json_dict_raw["tag"]:
+            sample["class_label"][class_index_dict[x]] = 1
+        sample["class_label"] = torch.tensor(sample["class_label"]).float()
     del sample[text_ext]
     sample["audio_name"] = sample["__key__"].split("/")[-1] + "." + audio_ext
     sample["text_name"] = sample["__key__"].split("/")[-1] + "." + text_ext
@@ -414,22 +640,40 @@ def preprocess(
     return sample
 
 
-# TODO: Yusong: clear unused arguments
+def collate_fn(batch):
+    """
+    Collate function for wdsdataloader.
+    batch: a list of dict, each dict is a sample
+    """
+    # concatenate values in each dictionary. if it is a tensor, concatenate. if it is a list, extend.
+    batch_dict = {}
+    for k in batch[0].keys():
+        if isinstance(batch[0][k], dict):  # dealwith bert tokenizer output
+            batch_dict[k] = {}
+            for kk in batch[0][k].keys():
+                tmp = []
+                for i in range(len(batch)):
+                    tmp.append(batch[i][k][kk])
+                batch_dict[k][kk] = torch.vstack(tmp)
+        elif isinstance(batch[0][k], torch.Tensor):
+            batch_dict[k] = torch.stack([sample[k] for sample in batch])
+        elif isinstance(batch[0][k], np.ndarray):
+            batch_dict[k] = torch.tensor(np.stack([sample[k] for sample in batch]))
+        else:
+            batch_dict[k] = [sample[k] for sample in batch]
+    return batch_dict
+
+
 def get_wds_dataset(
     args,
     model_cfg,
     is_train,
     audio_ext="flac",
     text_ext="json",
-    samplerate=48000,
-    mono=True,
     max_len=480000,
-    dtype="float64",
-    res_type="kaiser_best",
     proportion=1.0,
     sizefilepath_=None,
     is_local=None,
-    resample_method=None,
 ):
     """
     Get a dataset for wdsdataloader.
@@ -469,7 +713,8 @@ def get_wds_dataset(
 
     pipeline = [wds.SimpleShardList(input_shards)]
     # at this point we have an iterator over all the shards
-    if is_train:
+    # TODO: (yusong): add a if statement of distributed. If not, we don't need to split_by_node
+    if is_train or args.parallel_eval:
         pipeline.extend(
             [
                 wds.detshuffle(
@@ -497,38 +742,34 @@ def get_wds_dataset(
                 wds.tarfile_to_samples(handler=log_and_continue),
             ]
         )
-    pipeline.extend(
-        [
-            wds.map(
-                partial(
-                    preprocess,
-                    audio_ext=audio_ext,
-                    text_ext=text_ext,
-                    samplerate=samplerate,
-                    mono=mono,
-                    max_len=max_len,
-                    dtype=dtype,
-                    res_type=res_type,
-                    resample_method=args.resample_method,
-                )
-            ),
-            # TODO: (yusong) use wds.to_dict instead?
-            wds.to_tuple(
-                "__url__",
-                "__key__",
-                "waveform",
-                "text",
-                "raw_text",
-                "audio_name",
-                "text_name",
-                "audio_orig_sr",
-            ),
-            wds.batched(args.batch_size, partial=not is_train),
-        ]
+    pipeline.append(
+        wds.map(
+            partial(
+                preprocess,
+                audio_ext=audio_ext,
+                text_ext=text_ext,
+                max_len=max_len,
+                audio_cfg=model_cfg['audio_cfg'],
+                class_index_dict=copy.deepcopy(args.class_index_dict),
+                data_filling=args.data_filling,
+                data_truncating=args.data_truncating,
+                text_augment_selection=args.text_augment_selection,
+            )
+        ),
+    )
+
+    pipeline.append(
+        wds.batched(
+            args.batch_size,
+            partial=not (is_train or args.parallel_eval),
+            collation_fn=collate_fn,
+        )
     )
 
     dataset = wds.DataPipeline(*pipeline)
-    if is_train:
+    if is_train or args.parallel_eval:
+        # (yusong): Currently parallel evaluation will be not precise as we are repeat the last few samples.
+        # (yusong): See comments below.
         # roll over and repeat a few samples to get same number of full batches on each node
         global_batch_size = args.batch_size * args.world_size
         num_batches = math.ceil(num_samples / global_batch_size)
@@ -547,7 +788,7 @@ def get_wds_dataset(
 
     kwargs = {}
     if args.horovod:  # multi-node training on summit
-        kwargs['multiprocessing_context'] = 'forkserver'
+        kwargs["multiprocessing_context"] = "forkserver"
 
     dataloader = wds.WebLoader(
         dataset, batch_size=None, shuffle=False, num_workers=args.workers, **kwargs
@@ -675,26 +916,40 @@ def get_dataset_fn(data_path, dataset_type):
 
 
 def get_data(args, model_cfg):
-    # deprecated for audio
-    # preprocess_train, preprocess_val = preprocess_fns
     data = {}
 
-    # need to CHANGE when using the formal webdataset
-    # if args.train_data:
-    #     data["train"] = get_dataset_fn(args.train_data, args.dataset_type)(
-    #         args, preprocess_train, is_train=True
-    #     )
-
-    # if args.val_data:
-    #     data["val"] = get_dataset_fn(args.val_data, args.dataset_type)(
-    #         args, preprocess_val, is_train=False
-    #     )
+    args.class_index_dict = load_class_label(args.class_label_path)
+
+    if args.datasetinfos is None:
+        args.datasetinfos = ["train", "unbalanced_train", "balanced_train"]
+    if args.dataset_type == "webdataset":
+        args.train_data = get_tar_path_from_dataset_name(
+            args.datasetnames,
+            args.datasetinfos,
+            islocal=not args.remotedata,
+            proportion=args.dataset_proportion,
+            dataset_path=args.datasetpath,
+            full_dataset=args.full_train_dataset,
+        )
 
-    # if args.imagenet_val is not None:
-    #     data["imagenet-val"] = get_imagenet(args, preprocess_fns, "val")
+        if args.full_train_dataset is None:
+            args.full_train_dataset = []
+        if args.exclude_eval_dataset is None:
+            args.exclude_eval_dataset = []
+        excluded_eval_datasets = args.full_train_dataset + args.exclude_eval_dataset
+
+        val_dataset_names = [n for n in args.datasetnames if n not in excluded_eval_datasets] \
+            if excluded_eval_datasets else args.datasetnames
+        args.val_dataset_names = val_dataset_names
+        args.val_data = get_tar_path_from_dataset_name(
+            val_dataset_names,
+            ["valid", "test", "eval"],
+            islocal=not args.remotedata,
+            proportion=1,
+            dataset_path=args.datasetpath,
+            full_dataset=None,
+        )
 
-    # if args.imagenet_v2 is not None:
-    #     data["imagenet-v2"] = get_imagenet(args, preprocess_fns, "v2")
     if args.train_data:
         data["train"] = get_dataset_fn(args.train_data, args.dataset_type)(
             args, model_cfg, is_train=True
diff --git a/src/training/distributed.py b/src/training/distributed.py
index 8ae6fb2..adb0e92 100644
--- a/src/training/distributed.py
+++ b/src/training/distributed.py
@@ -1,6 +1,7 @@
 import os
 
 import torch
+import socket
 
 try:
     import horovod.torch as hvd
@@ -82,7 +83,9 @@ def init_distributed_device(args):
         os.environ['LOCAL_RANK'] = str(args.local_rank)
         os.environ['RANK'] = str(args.rank)
         os.environ['WORLD_SIZE'] = str(args.world_size)
-        print("Distributed training: local_rank={}, rank={}, world_size={}".format(args.local_rank, args.rank, args.world_size))
+        print(f"Distributed training: local_rank={args.local_rank}, "
+              f"rank={args.rank}, world_size={args.world_size}, "
+              f"hostname={socket.gethostname()}, pid={os.getpid()}")
     elif is_using_distributed():
         if 'SLURM_PROCID' in os.environ:
             # DDP via SLURM
@@ -97,6 +100,19 @@ def init_distributed_device(args):
                 world_size=args.world_size,
                 rank=args.rank,
             )
+        elif 'OMPI_COMM_WORLD_SIZE' in os.environ: # using Summit cluster
+            world_size = int(os.environ['OMPI_COMM_WORLD_SIZE'])
+            world_rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
+            local_rank = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
+            args.local_rank = local_rank
+            args.rank = world_rank
+            args.world_size = world_size
+            torch.distributed.init_process_group(
+                backend=args.dist_backend,
+                init_method=args.dist_url,
+                world_size=args.world_size,
+                rank=args.rank,
+            )
         else:
             # DDP via torchrun, torch.distributed.launch
             args.local_rank, _, _ = world_info_from_env()
@@ -106,6 +122,9 @@ def init_distributed_device(args):
             args.world_size = torch.distributed.get_world_size()
             args.rank = torch.distributed.get_rank()
         args.distributed = True
+        print(f"Distributed training: local_rank={args.local_rank}, "
+              f"rank={args.rank}, world_size={args.world_size}, "
+              f"hostname={socket.gethostname()}, pid={os.getpid()}")
 
     if torch.cuda.is_available():
         if args.distributed and not args.no_set_device_rank:
diff --git a/src/training/lp_main.py b/src/training/lp_main.py
new file mode 100644
index 0000000..96c7ed3
--- /dev/null
+++ b/src/training/lp_main.py
@@ -0,0 +1,651 @@
+from cmath import cos
+from inspect import getargs
+import logging
+import os
+import random
+from datetime import datetime
+import bisect
+import copy
+from sched import scheduler
+import numpy as np
+import torch
+import torch.backends.cudnn as cudnn
+from torch import optim
+from torch.cuda.amp import GradScaler
+import faulthandler
+import pathlib
+import argparse
+import time
+
+try:
+    import wandb
+except ImportError:
+    wandb = None
+
+try:
+    import torch.utils.tensorboard as tensorboard
+except ImportError:
+    tensorboard = None
+
+try:
+    import horovod.torch as hvd
+except ImportError:
+    hvd = None
+
+from open_clip import create_model_and_transforms, trace_model, create_model
+from training.data import get_data
+from training.params import parse_args
+from training.distributed import is_master, init_distributed_device, world_info_from_env
+from training.logger import setup_logging
+from training.scheduler import cosine_lr
+from training.lp_train import train_one_epoch, evaluate
+from open_clip.utils import get_tar_path_from_dataset_name, dataset_split, get_optimizer
+from open_clip.utils import load_p, load_class_label
+from open_clip.linear_probe import LinearProbe
+
+
+def maintain_ckpts(args, startidx, all_idx_len):
+    for i in reversed(range(startidx, all_idx_len)):
+        if os.path.exists(os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt")):
+            os.rename(
+                os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt"),
+                os.path.join(args.checkpoint_path, f"epoch_top_{i+1}.pt"),
+            )
+    if os.path.exists(
+        os.path.join(args.checkpoint_path, f"epoch_top_{all_idx_len}.pt")
+    ):
+        os.remove(os.path.join(args.checkpoint_path, f"epoch_top_{all_idx_len}.pt"))
+    return
+
+
+def update_top_k_performance(
+    new_metrics_inputs, current_top_k_ckpt_metrics, args, ckpt, bignumbetter=True
+):
+    """
+    Record the top-k performance of the current epoch.
+    current_top_k_metrics is a dictionary of the form: {1: top_1_ckpt_measure, 2: top_2_ckpt_measure, ...}
+    """
+    if isinstance(new_metrics_inputs, (list, tuple)):
+        new_metrics_inputs = np.mean(new_metrics_inputs)
+        return update_top_k_performance(
+            new_metrics_inputs,
+            current_top_k_ckpt_metrics,
+            args=args,
+            ckpt=ckpt,
+            bignumbetter=bignumbetter,
+        )
+    elif isinstance(new_metrics_inputs, dict):
+        new_metrics_inputs = np.mean(list(new_metrics_inputs.values()))
+        return update_top_k_performance(
+            new_metrics_inputs,
+            current_top_k_ckpt_metrics,
+            args=args,
+            ckpt=ckpt,
+            bignumbetter=bignumbetter,
+        )
+    elif isinstance(new_metrics_inputs, (float, int)):
+        update_flag = {k: False for k in current_top_k_ckpt_metrics.keys()}
+        sorted_keys = sorted(current_top_k_ckpt_metrics.keys())
+        sorted_values = sorted(
+            current_top_k_ckpt_metrics.values(), reverse=bignumbetter
+        )
+        sorted_values_ = copy.deepcopy(sorted_values)
+        sorted_values.append(new_metrics_inputs)
+        sorted_values = sorted(sorted_values, reverse=bignumbetter)
+        sorted_values = sorted_values[:-1]
+
+        if sorted_values == sorted_values_:
+            return current_top_k_ckpt_metrics, new_metrics_inputs
+        else:
+            for i in range(len(sorted_keys)):
+                if current_top_k_ckpt_metrics[sorted_keys[i]] != sorted_values[i]:
+                    current_top_k_ckpt_metrics[sorted_keys[i]] = sorted_values[i]
+                    update_flag[sorted_keys[i]] = True
+            for i in range(len(update_flag)):
+                if update_flag[i]:
+                    maintain_ckpts(args, i, len(sorted_keys))
+                    torch.save(
+                        ckpt,
+                        os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt"),
+                    )
+                    break
+            return current_top_k_ckpt_metrics, new_metrics_inputs
+
+
+# def updateifNone(a, b):
+#     a = b if None else a
+#     return a
+
+
+def is_pretrained_params(n):
+    return (
+        n.startswith("clap_model.transformer")
+        or n in ["clap_model.positional_embedding", "clap_model.text_projection"]
+        or n.startswith("clap_model.token_embedding")
+        or n.startswith("clap_model.ln_final")
+        or n.startswith("clap_model.logit_scale_t")
+    )
+
+
+def random_seed(seed=42, rank=0):
+    torch.manual_seed(seed + rank)
+    np.random.seed(seed + rank)
+    random.seed(seed + rank)
+
+def config_lp_optimizer(model, data, args):
+    # set wd-related params to 0 if use adam optimizer
+    if args.optimizer == "adam":
+        args.wd = 0
+        args.wd_pretrained = 0
+        args.wd_new = 0
+
+    in_clap = (
+        lambda n, p: n.startswith("clap_model")
+    )
+
+    named_parameters = list(model.named_parameters())
+
+    optimizer = {}
+    scheduler = {}
+
+    # freeze text encoder
+    text_freeze_parameters = [
+        p
+        for n, p in named_parameters
+        if n.startswith("clap_model.transformer")
+        or n in ["clap_model.positional_embedding", "clap_model.text_projection"]
+        or n.startswith("clap_model.token_embedding")
+        or n.startswith("clap_model.ln_final")
+    ]
+
+    if args.freeze_text:
+        logging.info("Freeze Text!!!!")
+        for k in text_freeze_parameters:
+            k.requires_grad = False
+
+    if not args.lp_freeze:
+        exclude = (
+            lambda n, p: p.ndim < 2
+            or "bn" in n
+            or "ln" in n
+            or "bias" in n
+            or "logit_scale" in n
+        )
+        include = lambda n, p: not exclude(n, p)
+
+        # (yusong): we do not split the learning rate anymore
+        # p for n, p in named_parameters if in_clap(n,p) and exclude(n, p) and p.requires_grad
+        gain_or_bias_params = [
+            p for n, p in named_parameters if exclude(n, p) and p.requires_grad
+        ]
+        # rest_params = [p for n, p in named_parameters if in_clap(n,p) and include(n, p) and p.requires_grad]
+        rest_params = [p for n, p in named_parameters if include(n, p) and p.requires_grad]
+
+        if args.train_data is None:
+            optimizer = None
+            scheduler = None
+        else:
+            total_steps = data["train"].dataloader.num_batches * args.epochs
+
+            if args.split_opt:
+                for x in ["lr", "beta1", "beta2", "eps", "wd"]:
+                    for y in ["_new", "_pretrained"]:
+                        if getattr(args, x + y) is None:
+                            setattr(args, x + y, getattr(args, x))
+
+                gain_or_bias_pretrained_params = [
+                    p
+                    for n, p in named_parameters
+                    if (exclude(n, p) and p.requires_grad) and is_pretrained_params(n)
+                ]
+                rest_pretrained_params = [
+                    p
+                    for n, p in named_parameters
+                    if (include(n, p) and p.requires_grad) and is_pretrained_params(n)
+                ]
+                gain_or_bias_new_params = [
+                    p
+                    for n, p in named_parameters
+                    if (exclude(n, p) and p.requires_grad) and (not is_pretrained_params(n))
+                ]
+                rest_new_params = [
+                    p
+                    for n, p in named_parameters
+                    if (include(n, p) and p.requires_grad) and (not is_pretrained_params(n))
+                ]
+
+                pretrained_params_optimizer = get_optimizer(
+                    [
+                        {"params": gain_or_bias_pretrained_params, "weight_decay": 0.0},
+                        {
+                            "params": rest_pretrained_params,
+                            "weight_decay": args.wd_pretrained,
+                        },
+                    ],
+                    lr=args.lr_pretrained,
+                    betas=(args.beta1_pretrained, args.beta2_pretrained),
+                    eps=args.eps_pretrained,
+                    momentum=args.momentum_pretrained,
+                    optimizer_name=args.optimizer,
+                )
+                pretrained_params_scheduler = cosine_lr(
+                    pretrained_params_optimizer,
+                    args.lr_pretrained,
+                    args.warmup,
+                    total_steps,
+                )
+
+                new_params_optimizer = get_optimizer(
+                    [
+                        {"params": gain_or_bias_new_params, "weight_decay": 0.0},
+                        {"params": rest_new_params, "weight_decay": args.wd_new},
+                    ],
+                    lr=args.lr_new,
+                    betas=(args.beta1_new, args.beta2_new),
+                    eps=args.eps_new,
+                    momentum=args.momentum_new,
+                    optimizer_name=args.optimizer,
+                )
+                new_params_scheduler = cosine_lr(
+                    new_params_optimizer, args.lr_new, args.warmup, total_steps
+                )
+
+                optimizer["text"] = pretrained_params_optimizer
+                optimizer["audio"] = new_params_optimizer
+                scheduler["text"] = pretrained_params_scheduler
+                scheduler["audio"] = new_params_scheduler
+
+                if args.horovod:
+                    pretrained_params_optimizer = hvd.DistributedOptimizer(
+                        pretrained_params_optimizer,
+                        named_parameters=model.named_parameters(),
+                    )
+                    new_params_optimizer = hvd.DistributedOptimizer(
+                        new_params_optimizer, named_parameters=model.named_parameters()
+                    )
+                    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
+                    hvd.broadcast_optimizer_state(pretrained_params_optimizer, root_rank=0)
+                    hvd.broadcast_optimizer_state(new_params_optimizer, root_rank=0)
+            else:
+
+                optimizer["clap"] = get_optimizer(
+                    [
+                        {"params": gain_or_bias_params, "weight_decay": 0.0},
+                        {"params": rest_params, "weight_decay": args.wd},
+                    ],
+                    lr=args.lr,
+                    betas=(args.beta1, args.beta2),
+                    eps=args.eps,
+                    momentum=args.momentum,
+                    optimizer_name=args.optimizer,
+            )
+                scheduler["clap"] = cosine_lr(optimizer["clap"], args.lr, args.warmup, total_steps)
+
+                if args.horovod:
+                    optimizer["clap"] = hvd.DistributedOptimizer(
+                        optimizer["clap"], named_parameters=model.named_parameters()
+                    )
+                    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
+                    hvd.broadcast_optimizer_state(optimizer["clap"], root_rank=0)
+
+    # linear probe optimizer
+    else:
+        lp_params = [p for n, p in named_parameters if (not in_clap(n, p)) and p.requires_grad]
+        lp_optim = get_optimizer(lp_params, lr=args.lp_lr, betas=(args.beta1, args.beta2), eps=args.eps, momentum=0.9,
+                                 optimizer_name=args.optimizer)
+        optimizer["lp"] = lp_optim
+
+    return optimizer, scheduler, text_freeze_parameters
+
+
+def main():
+    args = parse_args()
+
+    time.sleep(args.sleep)
+
+    # sanitize model name for filesystem / uri use, easier if we don't use / in name as a rule?
+    args.amodel = args.amodel.replace("/", "-")
+    # download sizes.json file
+
+    # (yusong): the below two lines are for debug
+    # print("setting up faulthandler")
+    # faulthandler.register(10)
+
+    random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed(args.seed)
+    torch.cuda.manual_seed_all(args.seed)
+    np.random.seed(args.seed)
+    args.class_index_dict = load_class_label(args.class_label_path)
+
+    # get the name of the experiments
+    if args.name is None:
+        args.name = "-".join(
+            [
+                datetime.now().strftime("%Y_%m_%d-%H_%M_%S"),
+                f"linear_probe"
+                f"model_{args.amodel}",
+                f"lr_{args.lr}",
+                f"b_{args.batch_size}",
+                f"j_{args.workers}",
+                f"p_{args.precision}",
+            ]
+        )
+
+    # discover initial world args early so we can log properly
+    args.distributed = False
+    args.local_rank, args.rank, args.world_size = world_info_from_env()
+
+    if args.remotedata and is_master(args):
+        for dataset_name in args.datasetnames:
+            for split in dataset_split[dataset_name]:
+                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
+                    os.makedirs(f"./json_files/{dataset_name}/{split}")
+                os.system(
+                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
+                )
+
+    args.log_path = None
+    if is_master(args, local=args.log_local):
+        log_base_path = os.path.join(args.logs, args.name)
+        os.makedirs(log_base_path, exist_ok=True)
+        log_filename = f"out-{args.rank}" if args.log_local else "out.log"
+        args.log_path = os.path.join(log_base_path, log_filename)
+
+        # avoid log dir in same name:
+        postfix = 0
+        while os.path.exists(args.log_path):
+            postfix += 1
+            log_base_path_new = log_base_path+'-'+str(postfix)
+            os.makedirs(log_base_path_new, exist_ok=True)
+            log_filename = f"out-{args.rank}" if args.log_local else "out.log"
+            args.log_path = os.path.join(log_base_path_new, log_filename)
+            # print(
+            #     "Error. Experiment already exists. Use --name {} to specify a new experiment."
+            # )
+            # return -1
+
+    # Set logger
+    args.log_level = logging.DEBUG if args.debug else logging.INFO
+    setup_logging(args.log_path, args.log_level)
+
+    # fully initialize distributed device environment
+    device = init_distributed_device(args)
+
+    args.wandb = "wandb" in args.report_to or "all" in args.report_to
+    args.tensorboard = "tensorboard" in args.report_to or "all" in args.report_to
+    if is_master(args):
+        args.tensorboard_path = (
+            os.path.join(args.logs, args.name, "tensorboard")
+            if args.tensorboard
+            else ""
+        )
+        args.checkpoint_path = os.path.join(args.logs, args.name, "checkpoints")
+        for dirname in [args.tensorboard_path, args.checkpoint_path]:
+            if dirname:
+                os.makedirs(dirname, exist_ok=True)
+    else:
+        args.tensorboard_path = ""
+        args.checkpoint_path = ""
+
+    if args.copy_codebase:
+        copy_codebase(args)
+
+    assert args.precision in ["amp", "fp16", "fp32"]
+    if args.precision == "fp16":
+        logging.warning(
+            "It is recommended to use AMP mixed-precision instead of FP16. "
+            "FP16 support needs further verification and tuning, especially for train."
+        )
+
+    if args.horovod:
+        logging.info(
+            f"Running in horovod mode with multiple processes / nodes. Device: {args.device}."
+            f"Process (global: {args.rank}, local {args.local_rank}), total {args.world_size}."
+        )
+    elif args.distributed:
+        logging.info(
+            f"Running in distributed mode with multiple processes. Device: {args.device}."
+            f"Process (global: {args.rank}, local {args.local_rank}), total {args.world_size}."
+        )
+    else:
+        logging.info(f"Running with a single process. Device {args.device}.")
+
+    logging.info(f'openai cache dir: {os.path.expanduser(args.openai_model_cache_dir)}')
+
+    # Create CLAP model
+    clap_model, clap_model_cfg = create_model(
+        args.amodel,
+        args.tmodel,
+        args.pretrained,
+        precision=args.precision,
+        device=device,
+        jit=args.torchscript,
+        force_quick_gelu=args.force_quick_gelu,
+        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
+        skip_params=False,
+        pretrained_audio=args.pretrained_audio,
+        pretrained_text=args.pretrained_text,
+        enable_fusion=args.enable_fusion,
+        fusion_type=args.fusion_type
+    )
+
+    args.lp_out_ch = len(list(args.class_index_dict.keys()))
+    # Linear Probe 
+    logging.info(f"linear probe using mlp: {args.lp_mlp}")
+    logging.info(f"linear probe using freeze: {args.lp_freeze}")
+    logging.info(f"linear probe act layer: {args.lp_act}")
+    logging.info(f"linear probe out ch: {args.lp_out_ch}")
+    logging.info(f"linear probe learning rate (if applicable): {args.lp_lr}")
+    logging.info(f"linear probe loss func: {args.lp_loss}")
+    logging.info(f"linear probe lp_metrics: {args.lp_metrics}")
+
+    model = LinearProbe(
+        clap_model,
+        mlp=args.lp_mlp, freeze=args.lp_freeze,
+        in_ch=512, out_ch=args.lp_out_ch,
+        act=args.lp_act
+    ) # in_ch is fixed (i.e., 512)
+    model = model.to(device)
+
+    if args.horovod:
+        with torch.no_grad():
+            for param in model.parameters():
+                param.set_(param.contiguous())
+
+    if args.trace:
+        model = trace_model(model, batch_size=args.batch_size, device=device)
+
+    if is_master(args):
+        logging.info("Linear Probe CLAP Model:")
+        logging.info(f"{str(clap_model)}")
+        logging.info("Params:")
+        params_file = os.path.join(args.logs, args.name, "params.txt")
+        with open(params_file, "w") as f:
+            for name in sorted(vars(args)):
+                val = getattr(args, name)
+                logging.info(f"  {name}: {val}")
+                f.write(f"{name}: {val}\n")
+
+    if args.distributed and not args.horovod:
+        if args.use_bn_sync:
+            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+        ddp_args = {}
+        if args.ddp_static_graph:
+            # this doesn't exist in older PyTorch, arg only added if enabled
+            ddp_args["static_graph"] = True
+        model = torch.nn.parallel.DistributedDataParallel(
+            model, device_ids=[device], find_unused_parameters=True, **ddp_args
+        )
+
+    data = get_data(args, clap_model_cfg)
+    assert len(data), "At least one train or eval dataset must be specified."
+    if args.trace:
+        assert "train" not in data, "Cannot train with traced model"
+
+
+    optimizer, scheduler, text_freeze_parameters = config_lp_optimizer(model, data, args)
+
+
+    scaler = GradScaler() if args.precision == "amp" else None
+
+    # optionally resume from a checkpoint
+    start_epoch = 0
+    if args.resume is not None:
+        if os.path.isfile(args.resume):
+            checkpoint = torch.load(args.resume, map_location=device)
+            if "epoch" in checkpoint:
+                # resuming a train checkpoint w/ epoch and optimizer state
+                start_epoch = checkpoint["epoch"]
+                sd = checkpoint["state_dict"]
+                if not args.distributed and next(iter(sd.items()))[0].startswith(
+                    "module"
+                ):
+                    sd = {k[len("module.") :]: v for k, v in sd.items()}
+                model.load_state_dict(sd)
+                if args.split_opt:
+                    if optimizer is not None:
+                        for k, o_ in optimizer.items():
+                            o_.load_state_dict(checkpoint[k + "_" + "optimizer"])
+                if optimizer is not None:
+                    optimizer.load_state_dict(checkpoint["optimizer"])
+                if scaler is not None and "scaler" in checkpoint:
+                    scaler.load_state_dict(checkpoint["scaler"])
+                logging.info(
+                    f"=> resuming checkpoint '{args.resume}' (epoch {start_epoch})"
+                )
+            else:
+                # loading a bare (model only) checkpoint for fine-tune or evaluation
+                model.load_state_dict(checkpoint)
+                logging.info(
+                    f"=> loaded checkpoint '{args.resume}' (epoch {start_epoch})"
+                )
+            if args.freeze_text:
+                print("Freeze Text!!!!")
+                for k in text_freeze_parameters:
+                    k.requires_grad = False
+        else:
+            logging.info("=> no checkpoint found at '{}'".format(args.resume))
+
+    cudnn.benchmark = True
+    cudnn.deterministic = False
+
+    # determine if this worker should save logs and checkpoints. only do so if it is rank == 0
+    args.save_logs = args.logs and args.logs.lower() != "none" and is_master(args)
+    writer = None
+    if args.save_logs and args.tensorboard:
+        assert tensorboard is not None, "Please install tensorboard."
+        writer = tensorboard.SummaryWriter(args.tensorboard_path)
+
+    if args.wandb and is_master(args):
+        assert wandb is not None, "Please install wandb."
+        logging.debug("Starting wandb.")
+        args.train_sz = data["train"].dataloader.num_samples
+        if args.val_data is not None:
+            args.val_sz = data["val"].dataloader.num_samples
+        # you will have to configure this for your project!
+        wandb.init(
+            project="clap",
+            notes=args.wandb_notes,
+            name=args.wandb_notes,
+            tags=[],
+            config=vars(args),
+        )
+        if args.debug:
+            wandb.watch(model, log="all")
+        wandb.save(params_file)
+        logging.debug("Finished loading wandb.")
+
+    if "train" not in data:
+        evaluate(model, data, start_epoch, args, writer)
+        return
+    elif start_epoch == 0 and "val" in data and not args.no_eval:
+        evaluate(model, data, 0, args, writer)
+    if args.save_top_performance:
+        current_top_k_ckpt_metrics = {
+            i: 0 for i in range(args.save_top_performance)
+        }  # initialize the top-k metric for ckpts to 0
+
+    for epoch in range(start_epoch, args.epochs):
+        # freeze the text param after (include) args.freeze_text_after, this is -1 by default
+        if epoch == args.freeze_text_after:
+            print("Text pretrained parameters are freezed since this epoch.")
+            for k in text_freeze_parameters:
+                k.requires_grad = False
+        if is_master(args):
+            logging.info(f"Start epoch {epoch}")
+
+        train_one_epoch(model, data, epoch, optimizer, scaler, scheduler, args, writer)
+        completed_epoch = epoch + 1
+
+        if any(v in data for v in ("val", "imagenet-val", "imagenet-v2")) and not args.no_eval:
+            metrics = evaluate(model, data, completed_epoch, args, writer)
+            if args.save_top_performance:
+                top_k_dataset = args.top_k_checkpoint_select_dataset
+                top_k_metric = args.top_k_checkpoint_select_metric
+                filtered_metrics = [
+                    v
+                    for k, v in metrics.items()
+                    if top_k_metric in k and top_k_dataset in k
+                ]  # check all R@10 metrics (all dataset) and use it to update the ckpt
+        # Saving checkpoints.
+        if args.save_logs:
+            opt_dict = {
+                    k + "_" + "optimizer": v.state_dict() for k, v in optimizer.items()
+                }
+            checkpoint_dict = {
+                "epoch": completed_epoch,
+                "name": args.name,
+                "state_dict": model.state_dict(),
+            }
+            checkpoint_dict.update(opt_dict)
+            if scaler is not None:
+                checkpoint_dict["scaler"] = scaler.state_dict()
+
+            if completed_epoch == args.epochs or (
+                args.save_frequency > 0 and (completed_epoch % args.save_frequency) == 0
+            ):
+                torch.save(
+                    checkpoint_dict,
+                    os.path.join(args.checkpoint_path, f"epoch_{completed_epoch}.pt"),
+                )
+            if args.save_most_recent:
+                torch.save(
+                    checkpoint_dict,
+                    os.path.join(args.checkpoint_path, f"epoch_latest.pt"),
+                )
+            if args.save_top_performance and not args.no_eval:
+                update_top_k_performance(
+                    filtered_metrics,
+                    current_top_k_ckpt_metrics,
+                    args,
+                    checkpoint_dict,
+                    bignumbetter=True,
+                )
+
+    if args.wandb and is_master(args):
+        wandb.finish()
+
+
+def copy_codebase(args):
+    from shutil import copytree, ignore_patterns
+
+    new_code_path = os.path.join(args.logs, args.name, "code")
+    if os.path.exists(new_code_path):
+        print(
+            f"Error. Experiment already exists at {new_code_path}. Use --name to specify a new experiment."
+        )
+        return -1
+    print(f"Copying codebase to {new_code_path}")
+    current_code_path = os.path.realpath(__file__)
+    for _ in range(3):
+        current_code_path = os.path.dirname(current_code_path)
+    copytree(
+        current_code_path, new_code_path, ignore=ignore_patterns("log", "logs", "wandb")
+    )
+    print("Done copying code.")
+    return 1
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/training/lp_train.py b/src/training/lp_train.py
new file mode 100644
index 0000000..f5a977f
--- /dev/null
+++ b/src/training/lp_train.py
@@ -0,0 +1,292 @@
+import json
+import logging
+import math
+import os
+import time
+from contextlib import suppress
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+try:
+    import wandb
+except ImportError:
+    wandb = None
+
+from open_clip import LPLoss, LPMetrics, lp_gather_features
+from open_clip.utils import do_mixup, get_mix_lambda
+from .distributed import is_master
+from .zero_shot import zero_shot_eval
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def unwrap_model(model):
+    if hasattr(model, "module"):
+        return model.module
+    else:
+        return model
+
+
+def train_one_epoch(
+        model, data, epoch, optimizer, scaler, scheduler, args, tb_writer=None, extra_suffix=""
+):
+    device = torch.device(args.device)
+    autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress
+    model.train()
+    loss = LPLoss(args.lp_loss)
+
+    dataloader, sampler = data["train"].dataloader, data["train"].sampler
+    if args.distributed and sampler is not None:
+        sampler.set_epoch(epoch)
+    num_batches_per_epoch = dataloader.num_batches
+    sample_digits = math.ceil(math.log(dataloader.num_samples + 1, 10))
+
+    # for toy dataset
+    if args.dataset_type == "toy":
+        dataloader.dataset.generate_queue()
+
+    loss_m = AverageMeter()
+    batch_time_m = AverageMeter()
+    data_time_m = AverageMeter()
+    end = time.time()
+
+    for i, batch in enumerate(dataloader):
+        step = num_batches_per_epoch * epoch + i
+
+        if isinstance(scheduler, dict):
+            for s in scheduler.values():
+                s(step)
+        else:
+            scheduler(step)
+
+        audio = batch # contains mel_spec, wavform, and longer list
+        class_label = batch['class_label']
+        # audio = audio.to(device=device, non_blocking=True)
+        class_label = class_label.to(device=device, non_blocking=True)
+
+        if args.mixup:
+            # https://github.com/RetroCirce/HTS-Audio-Transformer/blob/main/utils.py#L146
+            mix_lambda = torch.from_numpy(get_mix_lambda(0.5, len(audio["waveform"]))).to(device)
+            class_label = do_mixup(class_label, mix_lambda)
+        else:
+            mix_lambda = None
+
+        data_time_m.update(time.time() - end)
+        if isinstance(optimizer, dict):
+            for o_ in optimizer.values():
+                o_.zero_grad()
+        else:
+            optimizer.zero_grad()
+
+        with autocast():
+            pred = model(audio, mix_lambda=mix_lambda, device=device)
+            total_loss = loss(pred, class_label)
+
+        if isinstance(optimizer, dict):
+            if scaler is not None:
+                scaler.scale(total_loss).backward()
+                for o_ in optimizer.values():
+                    if args.horovod:
+                        o_.synchronize()
+                        scaler.unscale_(o_)
+                        with o_.skip_synchronize():
+                            scaler.step(o_)
+                    else:
+                        scaler.step(o_)
+                scaler.update()
+            else:
+                total_loss.backward()
+                for o_ in optimizer.values():
+                    o_.step()
+        else:
+            if scaler is not None:
+                scaler.scale(total_loss).backward()
+                if args.horovod:
+                    optimizer.synchronize()
+                    scaler.unscale_(optimizer)
+                    with optimizer.skip_synchronize():
+                        scaler.step(optimizer)
+                else:
+                    scaler.step(optimizer)
+                scaler.update()
+            else:
+                total_loss.backward()
+                optimizer.step()
+
+        # Note: we clamp to 4.6052 = ln(100), as in the original paper.
+        with torch.no_grad():
+            unwrap_model(model).clap_model.logit_scale_a.clamp_(0, math.log(100))
+            unwrap_model(model).clap_model.logit_scale_t.clamp_(0, math.log(100))
+
+        batch_time_m.update(time.time() - end)
+        end = time.time()
+        batch_count = i + 1
+
+        if is_master(args) and (i % 100 == 0 or batch_count == num_batches_per_epoch):
+            if isinstance(audio, dict):
+                batch_size = len(audio["waveform"])
+            else:
+                batch_size = len(audio)
+            num_samples = batch_count * batch_size * args.world_size
+            samples_per_epoch = dataloader.num_samples
+            percent_complete = 100.0 * batch_count / num_batches_per_epoch
+
+            # NOTE loss is coarsely sampled, just master node and per log update
+            loss_m.update(total_loss.item(), batch_size)
+            if isinstance(optimizer, dict):
+                logging.info(
+                    f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
+                    f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
+                    f"Data (t): {data_time_m.avg:.3f} "
+                    f"Batch (t): {batch_time_m.avg:.3f} "
+                    f"LR: {[o_.param_groups[0]['lr'] for o_ in optimizer.values()]}"
+                )
+                log_data = {
+                    "loss": loss_m.val,
+                    "data_time": data_time_m.val,
+                    "batch_time": batch_time_m.val,
+                    "lr": [o_.param_groups[0]["lr"] for o_ in optimizer.values()],
+                }
+            else:
+                logging.info(
+                    f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
+                    f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
+                    f"Data (t): {data_time_m.avg:.3f} "
+                    f"Batch (t): {batch_time_m.avg:.3f} "
+                    f"LR: {optimizer.param_groups[0]['lr']:5f} "
+                )
+
+                # Save train loss / etc. Using non avg meter values as loggers have their own smoothing
+                log_data = {
+                    "loss": loss_m.val,
+                    "data_time": data_time_m.val,
+                    "batch_time": batch_time_m.val,
+                    "lr": optimizer.param_groups[0]["lr"],
+                }
+            for name, val in log_data.items():
+                name = f"train{extra_suffix}/{name}"
+                if tb_writer is not None:
+                    tb_writer.add_scalar(name, val, step)
+                if args.wandb:
+                    assert wandb is not None, "Please install wandb."
+                    wandb.log({name: val, "step": step})
+
+            # resetting batch / data time meters per log window
+            batch_time_m.reset()
+            data_time_m.reset()
+    # end for
+
+def evaluate(model, data, epoch, args, tb_writer=None, extra_suffix=""):
+    metrics = {}
+    if not args.parallel_eval:
+        if not is_master(args):
+            return metrics
+    device = torch.device(args.device)
+    model.eval()
+
+    # CHANGE
+    # zero_shot_metrics = zero_shot_eval(model, data, epoch, args)
+    # metrics.update(zero_shot_metrics)
+    if is_master(args):
+        print('Evaluating...')
+        metric_names = args.lp_metrics.split(',')
+        eval_tool = LPMetrics(metric_names=metric_names)
+
+    autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress        
+    if "val" in data and (
+            args.val_frequency
+            and ((epoch % args.val_frequency) == 0 or epoch == args.epochs)
+    ):
+        if args.parallel_eval:
+            dataloader, sampler = data["val"].dataloader, data["val"].sampler
+            if args.distributed and sampler is not None:
+                sampler.set_epoch(epoch)
+            samples_per_val = dataloader.num_samples
+        else:
+            dataloader = data["val"].dataloader
+            num_samples = 0
+            samples_per_val = dataloader.num_samples
+            
+        eval_info = {
+            'pred': [],
+            'target': []
+        }
+        with torch.no_grad():
+            for i, batch in enumerate(dataloader):
+                audio = batch # contains mel_spec, wavform, and longer list
+                class_label = batch['class_label']
+                           
+                # audio = audio.to(device=device, non_blocking=True)
+                class_label = class_label.to(device=device, non_blocking=True)
+
+                with autocast():
+                    pred = model(audio, device=device)
+                    if args.parallel_eval:
+                        pred, class_label = lp_gather_features(pred, class_label, args.world_size, args.horovod)
+                    eval_info['pred'].append(pred)
+                    eval_info['target'].append(class_label)
+                    
+                num_samples += class_label.shape[0]
+
+                if (i % 100) == 0:  # and i != 0:
+                    logging.info(
+                        f"Eval Epoch: {epoch} [{num_samples} / {samples_per_val}]"
+                    )
+                    
+            if is_master(args):
+                eval_info['pred'] = torch.cat(eval_info['pred'], 0).cpu()
+                eval_info['target'] = torch.cat(eval_info['target'], 0).cpu()
+                metric_dict = eval_tool.evaluate_mertics(eval_info['pred'], eval_info['target'])
+                metrics.update(metric_dict)
+                if "epoch" not in metrics.keys():
+                    metrics.update({"epoch": epoch})
+
+    if is_master(args):
+        if not metrics:
+            return metrics
+
+        logging.info(
+            f"Eval Epoch: {epoch} "
+            + "\n".join(
+                [
+                    "\t".join([f"{m}: {round(metrics[m], 4):.4f}" ])
+                    for m in metrics
+                ]
+            )
+        )
+        if args.save_logs:
+            for name, val in metrics.items():
+                if tb_writer is not None:
+                    tb_writer.add_scalar(f"val{extra_suffix}/{name}", val, epoch)
+
+            with open(os.path.join(args.checkpoint_path, "results.jsonl"), "a+") as f:
+                f.write(json.dumps(metrics))
+                f.write("\n")
+
+        if args.wandb:
+            assert wandb is not None, "Please install wandb."
+            for name, val in metrics.items():
+                wandb.log({f"val{extra_suffix}/{name}": val, "epoch": epoch})
+
+        return metrics
+    else:
+        return metrics
diff --git a/src/training/main.py b/src/training/main.py
index 886a833..fefc13b 100644
--- a/src/training/main.py
+++ b/src/training/main.py
@@ -10,6 +10,8 @@
 import torch.backends.cudnn as cudnn
 from torch import optim
 from torch.cuda.amp import GradScaler
+import faulthandler
+import pathlib
 
 try:
     import wandb
@@ -33,7 +35,7 @@
 from training.params import parse_args
 from training.scheduler import cosine_lr
 from training.train import train_one_epoch, evaluate
-from open_clip.utils import get_tar_path_from_dataset_name, dataset_split
+from open_clip.utils import dataset_split, get_optimizer
 
 
 def maintain_ckpts(args, startidx, all_idx_len):
@@ -128,48 +130,29 @@ def random_seed(seed=42, rank=0):
 def main():
     args = parse_args()
     # sanitize model name for filesystem / uri use, easier if we don't use / in name as a rule?
-    args.model = args.model.replace("/", "-")
+    args.amodel = args.amodel.replace("/", "-")
     # download sizes.json file
 
+    # (yusong): the below two lines are for debug
+    # print("setting up faulthandler")
+    # faulthandler.register(10)
+
     random.seed(args.seed)
     torch.manual_seed(args.seed)
     torch.cuda.manual_seed(args.seed)
     torch.cuda.manual_seed_all(args.seed)
     np.random.seed(args.seed)
+    if args.tmodel == "bert" or args.tmodel == "roberta" or args.tmodel == "bart":
+        assert (
+            args.pretrained == "" or args.pretrained is None
+        ), "bert/roberta/bart text encoder does not support pretrained models."
 
-    if args.remotedata:
-        for dataset_name in args.datasetnames:
-            for split in dataset_split[dataset_name]:
-                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
-                    os.makedirs(f"./json_files/{dataset_name}/{split}")
-                os.system(
-                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
-                )
-
-    if args.datasetinfos is None:
-        args.datasetinfos = ["train", "unbalanced_train", "balanced_train"]
-    if args.dataset_type == "webdataset":
-        args.train_data = get_tar_path_from_dataset_name(
-            args.datasetnames,
-            args.datasetinfos,
-            islocal=not args.remotedata,
-            proportion=args.dataset_proportion,
-            dataset_path=args.datasetpath,
-        )
-        args.val_data = get_tar_path_from_dataset_name(
-            args.datasetnames,
-            ["valid", "test", "eval"],
-            islocal=not args.remotedata,
-            proportion=1,
-            dataset_path=args.datasetpath,
-        )
-        # args.val_data = get_tar_path_from_dataset_name(args.datasetnames, ["valid"], islocal=not args.remotedata, template=args.data_txt_example)
     # get the name of the experiments
     if args.name is None:
         args.name = "-".join(
             [
                 datetime.now().strftime("%Y_%m_%d-%H_%M_%S"),
-                f"model_{args.model}",
+                f"model_{args.amodel}",
                 f"lr_{args.lr}",
                 f"b_{args.batch_size}",
                 f"j_{args.workers}",
@@ -181,6 +164,15 @@ def main():
     args.distributed = False
     args.local_rank, args.rank, args.world_size = world_info_from_env()
 
+    if args.remotedata and is_master(args):
+        for dataset_name in args.datasetnames:
+            for split in dataset_split[dataset_name]:
+                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
+                    os.makedirs(f"./json_files/{dataset_name}/{split}")
+                os.system(
+                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
+                )
+
     args.log_path = None
     if is_master(args, local=args.log_local):
         log_base_path = os.path.join(args.logs, args.name)
@@ -239,16 +231,22 @@ def main():
     else:
         logging.info(f"Running with a single process. Device {args.device}.")
 
-    logging.info(f'openai cache dir: {os.path.expanduser(args.openai_model_cache_dir)}')
+    logging.info(f"openai cache dir: {os.path.expanduser(args.openai_model_cache_dir)}")
 
     model, model_cfg = create_model(
-        args.model,
+        args.amodel,
+        args.tmodel,
         args.pretrained,
         precision=args.precision,
         device=device,
         jit=args.torchscript,
         force_quick_gelu=args.force_quick_gelu,
         openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
+        skip_params=True,
+        pretrained_audio=args.pretrained_audio,
+        pretrained_text=args.pretrained_text,
+        enable_fusion=args.enable_fusion,
+        fusion_type=args.fusion_type
     )
 
     if args.horovod:
@@ -301,10 +299,7 @@ def main():
     text_freeze_parameters = [
         p
         for n, p in named_parameters
-        if n.startswith("transformer")
-        or n in ["positional_embedding", "text_projection"]
-        or n.startswith("token_embedding")
-        or n.startswith("ln_final")
+        if 'text_branch' in n
     ]
 
     if args.freeze_text:
@@ -317,6 +312,12 @@ def main():
     ]
     rest_params = [p for n, p in named_parameters if include(n, p) and p.requires_grad]
 
+    # set wd-related params to 0 if use adam optimizer
+    if args.optimizer == "adam":
+        args.wd = 0
+        args.wd_pretrained = 0
+        args.wd_new = 0
+
     if args.train_data is None:
         optimizer = None
         scheduler = None
@@ -349,8 +350,7 @@ def main():
                 for n, p in named_parameters
                 if (include(n, p) and p.requires_grad) and (not is_pretrained_params(n))
             ]
-
-            pretrained_params_optimizer = optim.AdamW(
+            pretrained_params_optimizer = get_optimizer(
                 [
                     {"params": gain_or_bias_pretrained_params, "weight_decay": 0.0},
                     {
@@ -361,15 +361,16 @@ def main():
                 lr=args.lr_pretrained,
                 betas=(args.beta1_pretrained, args.beta2_pretrained),
                 eps=args.eps_pretrained,
-            )
+                momentum=args.momentum_pretrained,
+                optimizer_name=args.optimizer,
+                )
             pretrained_params_scheduler = cosine_lr(
                 pretrained_params_optimizer,
                 args.lr_pretrained,
                 args.warmup,
                 total_steps,
             )
-
-            new_params_optimizer = optim.AdamW(
+            new_params_optimizer = get_optimizer(
                 [
                     {"params": gain_or_bias_new_params, "weight_decay": 0.0},
                     {"params": rest_new_params, "weight_decay": args.wd_new},
@@ -377,7 +378,10 @@ def main():
                 lr=args.lr_new,
                 betas=(args.beta1_new, args.beta2_new),
                 eps=args.eps_new,
-            )
+                momentum=args.momentum_new,
+                optimizer_name=args.optimizer,
+                )
+
             new_params_scheduler = cosine_lr(
                 new_params_optimizer, args.lr_new, args.warmup, total_steps
             )
@@ -403,7 +407,7 @@ def main():
                 hvd.broadcast_optimizer_state(pretrained_params_optimizer, root_rank=0)
                 hvd.broadcast_optimizer_state(new_params_optimizer, root_rank=0)
         else:
-            optimizer = optim.AdamW(
+            optimizer = get_optimizer(
                 [
                     {"params": gain_or_bias_params, "weight_decay": 0.0},
                     {"params": rest_params, "weight_decay": args.wd},
@@ -411,7 +415,10 @@ def main():
                 lr=args.lr,
                 betas=(args.beta1, args.beta2),
                 eps=args.eps,
+                momentum=args.momentum,
+                optimizer_name=args.optimizer,
             )
+
             scheduler = cosine_lr(optimizer, args.lr, args.warmup, total_steps)
 
             if args.horovod:
@@ -481,6 +488,7 @@ def main():
         wandb.init(
             project="clap",
             notes=args.wandb_notes,
+            name=args.wandb_notes,
             tags=[],
             config=vars(args),
         )
@@ -491,17 +499,16 @@ def main():
 
     if "train" not in data:
         evaluate(model, data, start_epoch, args, writer)
-        print('Start First Evaluation')
         return
-    elif start_epoch == 0 and "val" in data:
+    elif start_epoch == 0 and "val" in data and not args.no_eval:
         evaluate(model, data, 0, args, writer)
-        # pass
+        #  print(f'rank {args.rank}, Start First Evaluation')#  (yusong): for debug
     if args.save_top_performance:
         current_top_k_ckpt_metrics = {
             i: 0 for i in range(args.save_top_performance)
         }  # initialize the top-k metric for ckpts to 0
 
-    print('Start Training')
+    #  print(f'rank {args.rank}, Start Training') #  (yusong): for debug
     for epoch in range(start_epoch, args.epochs):
         # freeze the text param after (include) args.freeze_text_after, this is -1 by default
         if epoch == args.freeze_text_after:
@@ -514,7 +521,10 @@ def main():
         train_one_epoch(model, data, epoch, optimizer, scaler, scheduler, args, writer)
         completed_epoch = epoch + 1
 
-        if any(v in data for v in ("val", "imagenet-val", "imagenet-v2")):
+        if (
+            any(v in data for v in ("val", "imagenet-val", "imagenet-v2"))
+            and not args.no_eval
+        ):
             metrics = evaluate(model, data, completed_epoch, args, writer)
             if args.save_top_performance:
                 top_k_dataset = args.top_k_checkpoint_select_dataset
@@ -553,7 +563,7 @@ def main():
                     checkpoint_dict,
                     os.path.join(args.checkpoint_path, f"epoch_latest.pt"),
                 )
-            if args.save_top_performance:
+            if args.save_top_performance and not args.no_eval:
                 update_top_k_performance(
                     filtered_metrics,
                     current_top_k_ckpt_metrics,
diff --git a/src/training/main_.py b/src/training/main_.py
index 1495526..bf52a50 100644
--- a/src/training/main_.py
+++ b/src/training/main_.py
@@ -128,7 +128,7 @@ def random_seed(seed=42, rank=0):
 def main():
     args = parse_args()
     # sanitize model name for filesystem / uri use, easier if we don't use / in name as a rule?
-    args.model = args.model.replace("/", "-")
+    args.amodel = args.amodel.replace("/", "-")
     # download sizes.json file
     if args.remotedata:
         for dataset_name in args.datasetnames:
@@ -162,7 +162,7 @@ def main():
         args.name = "-".join(
             [
                 datetime.now().strftime("%Y_%m_%d-%H_%M_%S"),
-                f"model_{args.model}",
+                f"model_{args.amodel}",
                 f"lr_{args.lr}",
                 f"b_{args.batch_size}",
                 f"j_{args.workers}",
@@ -234,7 +234,7 @@ def main():
 
     # don't need the transform
     # model, preprocess_train, preprocess_val = create_model_and_transforms(
-    #     args.model,
+    #     args.amodel,
     #     args.pretrained,
     #     precision=args.precision,
     #     device=device,
@@ -243,7 +243,7 @@ def main():
     #     # pretrained_image=args.pretrained_image,
     # )
     model, model_cfg = create_model(
-        args.model,
+        args.amodel,
         args.pretrained,
         precision=args.precision,
         device=device,
diff --git a/src/training/mlaion.sh b/src/training/mlaion.sh
new file mode 100644
index 0000000..93ff650
--- /dev/null
+++ b/src/training/mlaion.sh
@@ -0,0 +1,76 @@
+#!/bin/bash
+#SBATCH --comment clap
+#SBATCH --partition=gpu
+#SBATCH --job-name=mclap
+#SBATCH --nodes 8
+#SBATCH --ntasks-per-node 8
+#SBATCH --cpus-per-gpu=6
+#SBATCH --gres=gpu:8
+#SBATCH --output=%x_%j.out
+#SBATCH --exclude gpu-st-p4d-24xlarge-[66,106,141,144,225,324,347-350]
+
+module load intelmpi
+source /opt/intel/mpi/latest/env/vars.sh
+export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-install/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+export NCCL_PROTO=simple
+export PATH=/opt/amazon/efa/bin:$PATH
+export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
+
+export FI_EFA_FORK_SAFE=1
+export FI_LOG_LEVEL=1
+export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
+
+#export NCCL_ALGO=ring
+export NCCL_DEBUG=info
+#export NCCL_DEBUG_SUBSYS=INIT,ENV,GRAPH,COLL
+
+export PYTHONFAULTHANDLER=1
+
+export CUDA_LAUNCH_BLOCKING=0
+export OMPI_MCA_mtl_base_verbose=1
+export FI_EFA_ENABLE_SHM_TRANSFER=0
+export FI_PROVIDER=efa
+export FI_EFA_TX_MIN_CREDITS=64
+export NCCL_TREE_THRESHOLD=0
+
+
+#export NCCL_P2P_DISABLE=1
+#export NCCL_IBEXT_DISABLE=1
+#export NCCL_SOCKET_IFNAME="eth0,en,eth,em,bond"
+
+# sent to sub script
+export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_PORT=12802
+export COUNT_NODE=`scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l`
+
+echo go $COUNT_NODE
+echo $HOSTNAMES
+
+source ./clap/bin/activate
+cd /home/tyz/CLAP/src
+
+srun --comment clap --cpu_bind=v --accel-bind=gn python -m training.main \
+    --save-frequency 50 \
+    --save-top-performance 3 \
+    --save-most-recent \
+    --dataset-type="webdataset" \
+    --precision="fp32" \
+    --pretrained="openai" \
+    --warmup 10000 \
+    --batch-size=184 \
+    --lr=1e-3 \
+    --wd=0.1 \
+    --epochs=400 \
+    --workers=10 \
+    --use-bn-sync \
+    --freeze-text \
+    --model PANN-14 \
+    --report-to "wandb" \
+    --wandb-notes "text-audio-freeze-text-lr-1e-3-8-dataset-model-pann-14" \
+    --datasetnames "Clotho" "audiocaps" "BBCSoundEffects" "audioset" "free_to_use_sounds" "paramount_motion" "sonniss_game_effects" "wesoundeffects" \
+    --datasetinfos "train" "unbalanced_train" "balanced_train" "valid" \
+    --top-k-checkpoint-select-dataset="Clotho-test" \
+    --top-k-checkpoint-select-metric="mAP@10" \
+    --seed 3407 \
+    --remotedata
diff --git a/src/training/params.py b/src/training/params.py
index e14df86..a527122 100644
--- a/src/training/params.py
+++ b/src/training/params.py
@@ -27,7 +27,7 @@ def parse_args():
     parser.add_argument(
         "--freeze-text",
         default=False,
-        action='store_true',
+        action="store_true",
         help="if you need to freeze the text encoder, make this True",
     )
     parser.add_argument(
@@ -64,25 +64,25 @@ def parse_args():
         "--dataset-type",
         choices=["webdataset", "csv", "auto", "toy"],
         default="auto",
-        help="Which type of dataset to process."
+        help="Which type of dataset to process.",
     )
     parser.add_argument(
         "--csv-separator",
         type=str,
         default="\t",
-        help="For csv-like datasets, which separator to use."
+        help="For csv-like datasets, which separator to use.",
     )
     parser.add_argument(
         "--csv-img-key",
         type=str,
         default="filepath",
-        help="For csv-like datasets, the name of the key for the image paths."
+        help="For csv-like datasets, the name of the key for the image paths.",
     )
     parser.add_argument(
         "--csv-caption-key",
         type=str,
         default="title",
-        help="For csv-like datasets, the name of the key for the captions."
+        help="For csv-like datasets, the name of the key for the captions.",
     )
     parser.add_argument(
         "--imagenet-val",
@@ -98,27 +98,50 @@ def parse_args():
     )
     parser.add_argument(
         "--datasetnames",
-        nargs='+',
+        nargs="+",
         default=None,
         help="If loading webdataset, spedify the dataset names to load. Can be some of these: Clotho, audioset, audiocaps, BBCSoundEffects",
     )
+    parser.add_argument(
+        "--full-train-dataset",
+        nargs="+",
+        default=None,
+        help="Which dataset will be trained with all the subsets. (train+test)",
+    )
+    parser.add_argument(
+        "--exclude-eval-dataset",
+        nargs="+",
+        default=None,
+        help="Which dataset will be excluded with evaluation",
+    )
     parser.add_argument(
         "--datasetinfos",
-        nargs='+',
+        nargs="+",
         default=None,
         help="If loading webdataset, spedify the dataset types to load. Can be some of these: train, test, valid, unbalanced_train, balanced_train, eval",
     )
-    parser.add_argument("--dataset-proportion", type=float, default=1.0, help="How much proportion of dataset we want to train.")
+    parser.add_argument(
+        "--dataset-proportion",
+        type=float,
+        default=1.0,
+        help="How much proportion of dataset we want to train.",
+    )
     parser.add_argument(
         "--remotedata",
         default=False,
-        action='store_true',
+        action="store_true",
         help="if the dataset is remote, set this flag",
     )
+    parser.add_argument(
+        "--class-label-path",
+        type=str,
+        default=None,
+        help="The path of the class label pickle or csv.",
+    )
     parser.add_argument(
         "--datasetpath",
         type=str,
-        default='/mnt/audio_clip/webdataset_tar',
+        default="/mnt/audio_clip/webdataset_tar",
         help="The path to the dataset",
     )
     parser.add_argument(
@@ -152,6 +175,7 @@ def parse_args():
     parser.add_argument("--beta1", type=float, default=None, help="Adam beta 1.")
     parser.add_argument("--beta2", type=float, default=None, help="Adam beta 2.")
     parser.add_argument("--eps", type=float, default=None, help="Adam epsilon.")
+    parser.add_argument("--momentum", type=float, default=None, help="SGD epsilon.")
     parser.add_argument("--wd", type=float, default=0.2, help="Weight decay.")
 
     parser.add_argument(
@@ -160,18 +184,42 @@ def parse_args():
         default=False,
         help="Use this flag to skip the learning rate decay.",
     )
-    parser.add_argument("--lr-pretrained", type=float, default=None, help="Learning rate for text.")
-    parser.add_argument("--beta1-pretrained", type=float, default=None, help="Adam beta 1 for text.")
-    parser.add_argument("--beta2-pretrained", type=float, default=None, help="Adam beta 2 for text.")
-    parser.add_argument("--eps-pretrained", type=float, default=None, help="Adam epsilon for text.")
-    parser.add_argument("--wd-pretrained", type=float, default=0.2, help="Weight decay for text.")
-
-    parser.add_argument("--lr-new", type=float, default=None, help="Learning rate for audio.")
-    parser.add_argument("--beta1-new", type=float, default=None, help="Adam beta 1 for audio.")
-    parser.add_argument("--beta2-new", type=float, default=None, help="Adam beta 2 for audio.")
-    parser.add_argument("--eps-new", type=float, default=None, help="Adam epsilon for audio.")
-    parser.add_argument("--wd-new", type=float, default=0.2, help="Weight decay for audio.")
-
+    parser.add_argument(
+        "--lr-pretrained", type=float, default=None, help="Learning rate for text."
+    )
+    parser.add_argument(
+        "--beta1-pretrained", type=float, default=None, help="Adam beta 1 for text."
+    )
+    parser.add_argument(
+        "--beta2-pretrained", type=float, default=None, help="Adam beta 2 for text."
+    )
+    parser.add_argument(
+        "--eps-pretrained", type=float, default=None, help="Adam epsilon for text."
+    )
+    parser.add_argument(
+        "--wd-pretrained", type=float, default=0.2, help="Weight decay for text."
+    )
+    parser.add_argument(
+        "--momentum-pretrained", type=float, default=0.9, help="Momentum for text."
+    )
+    parser.add_argument(
+        "--lr-new", type=float, default=None, help="Learning rate for audio."
+    )
+    parser.add_argument(
+        "--beta1-new", type=float, default=None, help="Adam beta 1 for audio."
+    )
+    parser.add_argument(
+        "--beta2-new", type=float, default=None, help="Adam beta 2 for audio."
+    )
+    parser.add_argument(
+        "--eps-new", type=float, default=None, help="Adam epsilon for audio."
+    )
+    parser.add_argument(
+        "--wd-new", type=float, default=0.2, help="Weight decay for audio."
+    )
+    parser.add_argument(
+        "--momentum-new", type=float, default=0.9, help="Momentum for audio."
+    )
     parser.add_argument(
         "--warmup", type=int, default=10000, help="Number of steps to warmup for."
     )
@@ -179,7 +227,8 @@ def parse_args():
         "--use-bn-sync",
         default=False,
         action="store_true",
-        help="Whether to use batch norm sync.")
+        help="Whether to use batch norm sync.",
+    )
     parser.add_argument(
         "--skip-scheduler",
         action="store_true",
@@ -190,7 +239,10 @@ def parse_args():
         "--save-frequency", type=int, default=1, help="How often to save checkpoints."
     )
     parser.add_argument(
-        "--save-top-performance", type=int, default=0, help="Save the top x performance weights if the value >0"
+        "--save-top-performance",
+        type=int,
+        default=0,
+        help="Save the top x performance weights if the value >0",
     )
     parser.add_argument(
         "--save-most-recent",
@@ -202,14 +254,10 @@ def parse_args():
         "--zeroshot-frequency", type=int, default=2, help="How often to run zero shot."
     )
     parser.add_argument(
-        "--val-frequency", type=int, default=1, help="How often to run evaluation with val data."
-    )
-    parser.add_argument(
-        "--resample-method",
-        default=None,
-        type=str,
-        required=True,
-        help="Set None to disable resampling. Set TorchAudio to use batchwise torchaudio.resample. Set Librosa to use librosa.resample in the dataloader. Set TorchAudio_ to use torchaudio.resample in the dataloader.",
+        "--val-frequency",
+        type=int,
+        default=1,
+        help="How often to run evaluation with val data.",
     )
     parser.add_argument(
         "--resume",
@@ -221,30 +269,48 @@ def parse_args():
         "--precision",
         choices=["amp", "fp16", "fp32"],
         default="amp",
-        help="Floating point precision."
+        help="Floating point precision.",
     )
     parser.add_argument(
-        "--model",
+        "--amodel",
         type=str,
         default="RN50",
-        help="Name of the vision backbone to use.",
+        help="Name of the audio backbone to use.",
+    )
+    parser.add_argument(
+        "--tmodel",
+        type=str,
+        default="transformer",
+        help="Name of the text backbone to use. Can be [transformer, bert, roberta, bart]",
+    )
+    parser.add_argument(
+        "--pretrained-audio",
+        default="",
+        type=str,
+        help="Use a pretrained audio model weights for the audio encoder of CLAP",
+    )
+    parser.add_argument(
+        "--pretrained-text",
+        default="",
+        type=str,
+        help="Use a pretrained text model weights for the text encoder of CLAP",
     )
     parser.add_argument(
         "--pretrained",
-        default='',
+        default="",
         type=str,
         help="Use a pretrained CLIP model weights with the specified tag or file path.",
     )
     parser.add_argument(
         "--pretrained-image",
         default=False,
-        action='store_true',
+        action="store_true",
         help="Load imagenet pretrained weights for image tower backbone if available.",
     )
     parser.add_argument(
         "--lock-image",
         default=False,
-        action='store_true',
+        action="store_true",
         help="Lock full image tower by disabling gradients.",
     )
     parser.add_argument(
@@ -256,37 +322,37 @@ def parse_args():
     parser.add_argument(
         "--lock-image-freeze-bn-stats",
         default=False,
-        action='store_true',
+        action="store_true",
         help="Freeze BatchNorm running stats in image tower for any locked layers.",
     )
     parser.add_argument(
         "--local-loss",
         default=False,
         action="store_true",
-        help="calculate loss w/ local features @ global (instead of realizing full global @ global matrix)"
+        help="calculate loss w/ local features @ global (instead of realizing full global @ global matrix)",
     )
     parser.add_argument(
         "--gather-with-grad",
         default=False,
         action="store_true",
-        help="enable full distributed gradient for feature gather"
+        help="enable full distributed gradient for feature gather",
     )
     parser.add_argument(
         "--force-quick-gelu",
         default=False,
-        action='store_true',
+        action="store_true",
         help="Force use of QuickGELU activation for non-OpenAI transformer models.",
     )
     parser.add_argument(
         "--torchscript",
         default=False,
-        action='store_true',
+        action="store_true",
         help="torch.jit.script the model, also uses jit version of OpenAI models if pretrained=='openai'",
     )
     parser.add_argument(
         "--trace",
         default=False,
-        action='store_true',
+        action="store_true",
         help="torch.jit.trace the model for inference / eval only",
     )
     # arguments for distributed training
@@ -301,15 +367,12 @@ def parse_args():
     )
     parser.add_argument(
         "--report-to",
-        default='',
+        default="",
         type=str,
-        help="Options are ['wandb', 'tensorboard', 'wandb,tensorboard']"
+        help="Options are ['wandb', 'tensorboard', 'wandb,tensorboard']",
     )
     parser.add_argument(
-        "--wandb-notes",
-        default='',
-        type=str,
-        help="Notes if logging with wandb"
+        "--wandb-notes", default="", type=str, help="Notes if logging with wandb"
     )
     parser.add_argument(
         "--C", type=float, default=3.16, help="inverse regularizer for logistic reg."
@@ -318,53 +381,181 @@ def parse_args():
         "--debug",
         default=False,
         action="store_true",
-        help="If true, more information is logged."
+        help="If true, more information is logged.",
     )
     parser.add_argument(
         "--copy-codebase",
         default=False,
         action="store_true",
-        help="If true, we copy the entire base on the log diretory, and execute from there."
+        help="If true, we copy the entire base on the log diretory, and execute from there.",
     )
     parser.add_argument(
         "--horovod",
         default=False,
         action="store_true",
-        help="Use horovod for distributed training."
+        help="Use horovod for distributed training.",
     )
     parser.add_argument(
         "--ddp-static-graph",
         default=False,
-        action='store_true',
+        action="store_true",
         help="Enable static graph optimization for DDP in PyTorch >= 1.11.",
     )
     parser.add_argument(
         "--no-set-device-rank",
         default=False,
         action="store_true",
-        help="Don't set device index from local rank (when CUDA_VISIBLE_DEVICES restricted to one per proc)."
+        help="Don't set device index from local rank (when CUDA_VISIBLE_DEVICES restricted to one per proc).",
     )
+    parser.add_argument("--seed", type=int, default=4242, help="Default random seed.")
+
     parser.add_argument(
-        "--seed", type=int, default=4242, help="Default random seed."
+        "--top-k-checkpoint-select-dataset",
+        type=str,
+        default="all",
+        help="The dataset of selecting top-k checkpoint.",
     )
 
+    # @R10, @R@5, @R1, mAP@10
     parser.add_argument(
-        "--top-k-checkpoint-select-dataset", type=str, default='all', help="The dataset of selecting top-k checkpoint."
+        "--top-k-checkpoint-select-metric",
+        type=str,
+        default="_R@10",
+        help="The metric for selecting top-k checkpoint.",
+    )
+    parser.add_argument(
+        "--openai-model-cache-dir",
+        type=str,
+        default="~/.cache/clip",
+        help="Directory to download OpenAI models.",
+    )
+    parser.add_argument(
+        "--optimizer",
+        type=str,
+        default="adamw",
+        help="can be AdamW or SGD",
+    )
+    parser.add_argument(
+        "--parallel-eval",
+        default=False,
+        action="store_true",
+        help="Eval in parallel (multi-GPU, multi-node).",
+    )
+
+    parser.add_argument(
+        "--no-eval",
+        default=False,
+        action="store_true",
+        help="Training without evaluation.",
+    )
+
+    parser.add_argument(
+        "--lp-mlp",
+        default=False,
+        action="store_true",
+        help="Linear Probe using MLP layer or not.",
+    )
+
+    parser.add_argument(
+        "--lp-freeze",
+        default=False,
+        action="store_true",
+        help="Linear Probe using Freeze CLAP or not",
+    )
+
+    parser.add_argument(
+        "--lp-act",
+        default="None",
+        type=str,
+        help="Options are ['relu','elu','prelu','softmax','sigmoid']",
+    )
+
+    parser.add_argument(
+        "--lp-loss", type=str, default="bce", help="Loss func of Linear Probe."
+    )
+
+    parser.add_argument(
+        "--lp-metrics",
+        type=str,
+        default="map,mauc,acc",
+        help="Metrics of Linear Probe.",
+    )
+
+    parser.add_argument(
+        "--lp-lr", type=float, default=1e-4, help="learning rate of linear probe"
+    )
+    parser.add_argument(
+        "--kappa", type=float, default=0,
+        help="the kappa in the weighted contrastive loss, default is to turn off the weighted contrastive loss"
+    )
+
+    parser.add_argument(
+        "--data-filling",
+        type=str,
+        default="pad",
+        help="type of data filling when the audio length is shorter than the max length."
+             "Can be one of the following: repeat, repeatpad, pad",
+    )
+    parser.add_argument(
+        "--data-truncating",
+        type=str,
+        default="rand_trunc",
+        help="type of data truncation when the audio length is longer than the max length."
+             "Can be one of the following: rand_trunc, fusion",
+    )
+
+    parser.add_argument(
+        "--clap-mlploss",
+        default=False,
+        action="store_true",
+        help="Using MLP loss for CLAP model or not",
+    )
+
+    parser.add_argument(
+        "--wandb-id",
+        type=str,
+        default=None,
+        help="the id of wandb experiment to restore.",
+    )
+
+    parser.add_argument(
+        "--sleep", type=float, default=0, help="sleep n seconds before start training"
+    )
+
+    # variable length processing
+    parser.add_argument(
+        "--enable-fusion",
+        default=False,
+        action="store_true",
+        help="Enable feature funsion for variable-length data",
     )
 
-    # @R10, @R@5, @R1, mAP@10
     parser.add_argument(
-        "--top-k-checkpoint-select-metric", type=str, default='_R@10', help="The metric for selecting top-k checkpoint."
+        "--fusion-type",
+        type=str,
+        default='None',
+        help="Type is among ['channel_map', 'daf_1d','aff_1d','iaff_1d','daf_2d','aff_2d','iaff_2d']",
     )
+
     parser.add_argument(
-        "--openai-model-cache-dir", type=str, default='~/.cache/clip', help="Directory to download OpenAI models."
+        "--mixup",
+        default=False,
+        action="store_true",
+        help="Enable mixup in finetuning training.",
     )
+    parser.add_argument(
+        "--text-augment-selection",
+        type=str,
+        default=None,
+        help="For selecting levels of augmented text. Type is among ['all', 'augment_only', 'none']",
+    )
+
     args = parser.parse_args()
 
     # If some params are not passed, we use the default values based on model name.
-    default_params = get_default_params(args.model)
+    default_params = get_default_params(args.amodel)
     for name, val in default_params.items():
         if getattr(args, name) is None:
             setattr(args, name, val)
 
-    return args
\ No newline at end of file
+    return args
diff --git a/src/training/train.py b/src/training/train.py
index 63502af..f3e41b1 100644
--- a/src/training/train.py
+++ b/src/training/train.py
@@ -14,7 +14,7 @@
 except ImportError:
     wandb = None
 
-from open_clip import ClipLoss
+from open_clip import ClipLoss, gather_features
 from .distributed import is_master
 from .zero_shot import zero_shot_eval
 
@@ -46,7 +46,7 @@ def unwrap_model(model):
 
 
 def train_one_epoch(
-    model, data, epoch, optimizer, scaler, scheduler, args, tb_writer=None
+        model, data, epoch, optimizer, scaler, scheduler, args, tb_writer=None
 ):
     device = torch.device(args.device)
     autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress
@@ -58,6 +58,8 @@ def train_one_epoch(
         rank=args.rank,
         world_size=args.world_size,
         use_horovod=args.horovod,
+        mlp_loss=args.clap_mlploss,
+        weight_loss_kappa=args.kappa,
     )
 
     dataloader, sampler = data["train"].dataloader, data["train"].sampler
@@ -76,16 +78,17 @@ def train_one_epoch(
     end = time.time()
 
     for i, batch in enumerate(dataloader):
+        # logging.info(f"batch {i} of {num_batches_per_epoch}")
         step = num_batches_per_epoch * epoch + i
         if isinstance(scheduler, dict):
             for s in scheduler.values():
                 s(step)
         else:
             scheduler(step)
-        audios = batch[2]  # (yusong) todo:  change to retrieve from index for now.
-        texts = batch[3][:, 0, :]
-        audios = audios.to(device=device, non_blocking=True)
-        texts = texts.to(device=device, non_blocking=True)
+        audios = batch  # contains mel_spec, wavform, and longer list
+        texts = batch['text']
+        # audios = audios.to(device=device, non_blocking=True)
+        # texts = texts.to(device=device, non_blocking=True)
 
         data_time_m.update(time.time() - end)
         if isinstance(optimizer, dict):
@@ -102,15 +105,23 @@ def train_one_epoch(
                 text_features_mlp,
                 logit_scale_a,
                 logit_scale_t,
-            ) = model(audios, texts)
-            total_loss = loss(
-                audio_features,
-                text_features,
-                audio_features_mlp,
-                text_features_mlp,
-                logit_scale_a,
-                logit_scale_t,
-            )
+            ) = model(audios, texts, device)
+
+            if args.clap_mlploss:
+                total_loss = loss(
+                    audio_features=audio_features,
+                    text_features=text_features,
+                    logit_scale_a=logit_scale_a,
+                    logit_scale_t=logit_scale_t,
+                    audio_features_mlp=audio_features_mlp,
+                    text_features_mlp=text_features_mlp
+                )
+            else:
+                total_loss = loss(
+                    audio_features=audio_features,
+                    text_features=text_features,
+                    logit_scale_a=logit_scale_a
+                )
         if isinstance(optimizer, dict):
             if scaler is not None:
                 scaler.scale(total_loss).backward()
@@ -145,13 +156,17 @@ def train_one_epoch(
         # Note: we clamp to 4.6052 = ln(100), as in the original paper.
         with torch.no_grad():
             unwrap_model(model).logit_scale_a.clamp_(0, math.log(100))
-            unwrap_model(model).logit_scale_t.clamp_(0, math.log(100))
+            if args.clap_mlploss:
+                unwrap_model(model).logit_scale_t.clamp_(0, math.log(100))
 
         batch_time_m.update(time.time() - end)
         end = time.time()
         batch_count = i + 1
         if is_master(args) and (i % 100 == 0 or batch_count == num_batches_per_epoch):
-            batch_size = len(audios)
+            if isinstance(audios, dict):
+                batch_size = len(audios["waveform"])
+            else:
+                batch_size = len(audios)
             num_samples = batch_count * batch_size * args.world_size
             samples_per_epoch = dataloader.num_samples
             percent_complete = 100.0 * batch_count / num_batches_per_epoch
@@ -161,43 +176,80 @@ def train_one_epoch(
             logit_scale_scalar_a = logit_scale_a.item()
             logit_scale_scalar_t = logit_scale_t.item()
             if isinstance(optimizer, dict):
-                logging.info(
-                    f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
-                    f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
-                    f"Data (t): {data_time_m.avg:.3f} "
-                    f"Batch (t): {batch_time_m.avg:.3f} "
-                    f"LR: {[o_.param_groups[0]['lr'] for o_ in optimizer.values()]} "
-                    f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
-                    f"Logit Scale Text: {logit_scale_scalar_t:.3f}"
-                )
-                log_data = {
-                    "loss": loss_m.val,
-                    "data_time": data_time_m.val,
-                    "batch_time": batch_time_m.val,
-                    "scale_audio": logit_scale_scalar_a,
-                    "scale_text": logit_scale_scalar_t,
-                    "lr": [o_.param_groups[0]["lr"] for o_ in optimizer.values()],
-                }
+                if args.clap_mlploss:
+                    logging.info(
+                        f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
+                        f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
+                        f"Data (t): {data_time_m.avg:.3f} "
+                        f"Batch (t): {batch_time_m.avg:.3f} "
+                        f"LR: {[o_.param_groups[0]['lr'] for o_ in optimizer.values()]} "
+                        f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
+                        f"Logit Scale Text: {logit_scale_scalar_t:.3f}"
+                    )
+                    log_data = {
+                        "loss": loss_m.val,
+                        "data_time": data_time_m.val,
+                        "batch_time": batch_time_m.val,
+                        "scale_audio": logit_scale_scalar_a,
+                        "scale_text": logit_scale_scalar_t,
+                        "lr": [o_.param_groups[0]["lr"] for o_ in optimizer.values()],
+                    }
+                else:
+                    logging.info(
+                        f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
+                        f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
+                        f"Data (t): {data_time_m.avg:.3f} "
+                        f"Batch (t): {batch_time_m.avg:.3f} "
+                        f"LR: {[o_.param_groups[0]['lr'] for o_ in optimizer.values()]} "
+                        f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
+                    )
+                    log_data = {
+                        "loss": loss_m.val,
+                        "data_time": data_time_m.val,
+                        "batch_time": batch_time_m.val,
+                        "scale_audio": logit_scale_scalar_a,
+                        "lr": [o_.param_groups[0]["lr"] for o_ in optimizer.values()],
+                    }
+
             else:
-                logging.info(
-                    f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
-                    f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
-                    f"Data (t): {data_time_m.avg:.3f} "
-                    f"Batch (t): {batch_time_m.avg:.3f} "
-                    f"LR: {optimizer.param_groups[0]['lr']:5f} "
-                    f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
-                    f"Logit Scale Text: {logit_scale_scalar_t:.3f}"
-                )
+                if args.clap_mlploss:
+                    logging.info(
+                        f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
+                        f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
+                        f"Data (t): {data_time_m.avg:.3f} "
+                        f"Batch (t): {batch_time_m.avg:.3f} "
+                        f"LR: {optimizer.param_groups[0]['lr']:5f} "
+                        f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
+                        f"Logit Scale Text: {logit_scale_scalar_t:.3f}"
+                    )
 
-                # Save train loss / etc. Using non avg meter values as loggers have their own smoothing
-                log_data = {
-                    "loss": loss_m.val,
-                    "data_time": data_time_m.val,
-                    "batch_time": batch_time_m.val,
-                    "scale_audio": logit_scale_scalar_a,
-                    "scale_text": logit_scale_scalar_t,
-                    "lr": optimizer.param_groups[0]["lr"],
-                }
+                    # Save train loss / etc. Using non avg meter values as loggers have their own smoothing
+                    log_data = {
+                        "loss": loss_m.val,
+                        "data_time": data_time_m.val,
+                        "batch_time": batch_time_m.val,
+                        "scale_audio": logit_scale_scalar_a,
+                        "scale_text": logit_scale_scalar_t,
+                        "lr": optimizer.param_groups[0]["lr"],
+                    }
+                else:
+                    logging.info(
+                        f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
+                        f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
+                        f"Data (t): {data_time_m.avg:.3f} "
+                        f"Batch (t): {batch_time_m.avg:.3f} "
+                        f"LR: {optimizer.param_groups[0]['lr']:5f} "
+                        f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
+                    )
+
+                    # Save train loss / etc. Using non avg meter values as loggers have their own smoothing
+                    log_data = {
+                        "loss": loss_m.val,
+                        "data_time": data_time_m.val,
+                        "batch_time": batch_time_m.val,
+                        "scale_audio": logit_scale_scalar_a,
+                        "lr": optimizer.param_groups[0]["lr"],
+                    }
             for name, val in log_data.items():
                 name = "train/" + name
                 if tb_writer is not None:
@@ -214,19 +266,33 @@ def train_one_epoch(
 
 def evaluate(model, data, epoch, args, tb_writer=None):
     metrics = {}
-    if not is_master(args):
-        return metrics
+    if not args.parallel_eval:
+        if not is_master(args):
+            return metrics
     device = torch.device(args.device)
     model.eval()
 
     # CHANGE
     # zero_shot_metrics = zero_shot_eval(model, data, epoch, args)
     # metrics.update(zero_shot_metrics)
-
+    if is_master(args):
+        print('Evaluating...')
     autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress
-    if "val" in data and (
-        args.val_frequency
-        and ((epoch % args.val_frequency) == 0 or epoch == args.epochs)
+    if args.val_dataset_names == ['Clotho', 'audiocaps']:
+        # if only clotho and audiocaps are used, then we will use a different evaluation function.
+        # This is because in the Clotho and audiocaps valid and test set, there are 5 text for 1 audio.
+        if args.parallel_eval:
+            # (yusong): just a hack here. Don't use parallel eval when evaluating only clotho and audiocaps.
+            raise NotImplementedError("Parallel evaluation not supported for eval only Clotho and audiocaps.")
+        val_metrics_per_dataset = evaluate_clotho_audiocaps(model, data, epoch, args, autocast, device, tb_writer)
+        for m in val_metrics_per_dataset.values():
+            metrics.update(m)
+        if "epoch" not in metrics.keys():
+            metrics.update({"epoch": epoch})
+        metrics = select_top_metric_clotho_audiocaps(metrics, val_metrics_per_dataset, args)
+    elif "val" in data and (
+            args.val_frequency
+            and ((epoch % args.val_frequency) == 0 or epoch == args.epochs)
     ):
         dataloader = data["val"].dataloader
         num_samples = 0
@@ -235,35 +301,48 @@ def evaluate(model, data, epoch, args, tb_writer=None):
         # FIXME this does not scale past small eval datasets
         # all_audio_features @ all_text_features will blow up memory and compute very quickly
         eval_info = {}
-        eval_info["all"] = {
-            "cumulative_loss": 0.0,
-            "num_samples": 0,
-            "all_audio_features": [],
-            "all_text_features": [],
-            "all_audio_features_mlp": [],
-            "all_text_features_mlp": [],
-        }  # cumulative_loss = 0.0
+        if args.clap_mlploss:
+            eval_info["all"] = {
+                "cumulative_loss": 0.0,
+                "num_samples": 0,
+                "all_audio_features": [],
+                "all_text_features": [],
+                "all_audio_features_mlp": [],
+                "all_text_features_mlp": []
+            }  # cumulative_loss = 0.0
+        else:
+            eval_info["all"] = {
+                "cumulative_loss": 0.0,
+                "num_samples": 0,
+                "all_audio_features": [],
+                "all_text_features": []
+            }  # cumu
         # all_audio_features, all_text_features, all_audio_features_mlp, all_text_features_mlp = [], [], [], []
         with torch.no_grad():
             for i, batch in enumerate(dataloader):
-                audios = batch[
-                    2
-                ]  # (yusong) todo:  change to retrieve from index for now.
-                texts = batch[3][:, 0, :]
-                audios = audios.to(device=device, non_blocking=True)
-                texts = texts.to(device=device, non_blocking=True)
-                all_names = list(set(["-".join(b.split("/")[-3:-1]) for b in batch[0]]))
+                audios = batch  # contains mel_spec, wavform, and longer list
+                texts = batch['text']
+                # audios = audios.to(device=device, non_blocking=True)
+
+                all_names = list(set(["-".join(b.split("/")[-3:-1]) for b in batch['__url__']]))
                 for name in all_names:
                     if name not in eval_info.keys():
-                        eval_info[name] = {
-                            "cumulative_loss": 0.0,
-                            "num_samples": 0,
-                            "all_audio_features": [],
-                            "all_text_features": [],
-                            "all_audio_features_mlp": [],
-                            "all_text_features_mlp": [],
-                        }
-
+                        if args.clap_mlploss:
+                            eval_info[name] = {
+                                "cumulative_loss": 0.0,
+                                "num_samples": 0,
+                                "all_audio_features": [],
+                                "all_text_features": [],
+                                "all_audio_features_mlp": [],
+                                "all_text_features_mlp": [],
+                            }
+                        else:
+                            eval_info[name] = {
+                                "cumulative_loss": 0.0,
+                                "num_samples": 0,
+                                "all_audio_features": [],
+                                "all_text_features": []
+                            }
                 with autocast():
                     (
                         audio_features,
@@ -272,149 +351,219 @@ def evaluate(model, data, epoch, args, tb_writer=None):
                         text_features_mlp,
                         logit_scale_a,
                         logit_scale_t,
-                    ) = model(audios, texts)
-
-                    num_samples += audio_features.shape[0]
-
-                    for n in [*all_names, "all"]:
-                        if n == "all":
-                            eval_info[n]["all_audio_features"].append(
-                                audio_features.cpu()
-                            )
-                            eval_info[n]["all_text_features"].append(
-                                text_features.cpu()
-                            )
-                            eval_info[n]["all_audio_features_mlp"].append(
-                                audio_features_mlp.cpu()
-                            )
-                            eval_info[n]["all_text_features_mlp"].append(
-                                text_features_mlp.cpu()
+                    ) = model(audios, texts, device)
+
+                    if args.parallel_eval:
+                        # multi-GPU eval
+                        if args.clap_mlploss:
+                            (
+                                audio_features,
+                                text_features,
+                                audio_features_mlp,
+                                text_features_mlp,
+                            ) = gather_features(
+                                audio_features=audio_features,
+                                text_features=text_features,
+                                audio_features_mlp=audio_features_mlp,
+                                text_features_mlp=text_features_mlp,
+                                local_loss=False,
+                                gather_with_grad=False,
+                                rank=args.rank,
+                                world_size=args.world_size,
+                                use_horovod=args.horovod,
+                                mlp_loss=args.clap_mlploss
                             )
                         else:
-                            idx = np.where(
-                                np.array(
-                                    ["-".join(b.split("/")[-3:-1]) for b in batch[0]]
-                                )
-                                == n
-                            )[0]
-                            eval_info[n]["all_audio_features"].append(
-                                audio_features.cpu().index_select(
-                                    0, torch.tensor(idx).long()
-                                )
+                            (
+                                audio_features,
+                                text_features,
+                            ) = gather_features(
+                                audio_features=audio_features,
+                                text_features=text_features,
+                                local_loss=False,
+                                gather_with_grad=False,
+                                rank=args.rank,
+                                world_size=args.world_size,
+                                use_horovod=args.horovod,
+                                mlp_loss=args.clap_mlploss
                             )
-                            eval_info[n]["all_text_features"].append(
-                                text_features.cpu().index_select(
-                                    0, torch.tensor(idx).long()
+
+                    if is_master(args):
+                        num_samples += audio_features.shape[0]
+                        for n in [*all_names, "all"]:
+                            if n == "all":
+                                eval_info[n]["all_audio_features"].append(
+                                    audio_features.cpu()
                                 )
-                            )
-                            eval_info[n]["all_audio_features_mlp"].append(
-                                audio_features_mlp.cpu().index_select(
-                                    0, torch.tensor(idx).long()
+                                eval_info[n]["all_text_features"].append(
+                                    text_features.cpu()
                                 )
-                            )
-                            eval_info[n]["all_text_features_mlp"].append(
-                                text_features_mlp.cpu().index_select(
-                                    0, torch.tensor(idx).long()
+                                if args.clap_mlploss:
+                                    eval_info[n]["all_audio_features_mlp"].append(
+                                        audio_features_mlp.cpu()
+                                    )
+                                    eval_info[n]["all_text_features_mlp"].append(
+                                        text_features_mlp.cpu()
+                                    )
+                            else:
+                                idx = np.where(
+                                    np.array(
+                                        ["-".join(b.split("/")[-3:-1]) for b in batch['__url__']]
+                                    )
+                                    == n
+                                )[0]
+                                eval_info[n]["all_audio_features"].append(
+                                    audio_features.cpu().index_select(
+                                        0, torch.tensor(idx).long()
+                                    )
                                 )
-                            )
+                                eval_info[n]["all_text_features"].append(
+                                    text_features.cpu().index_select(
+                                        0, torch.tensor(idx).long()
+                                    )
+                                )
+                                if args.clap_mlploss:
+                                    eval_info[n]["all_audio_features_mlp"].append(
+                                        audio_features_mlp.cpu().index_select(
+                                            0, torch.tensor(idx).long()
+                                        )
+                                    )
+                                    eval_info[n]["all_text_features_mlp"].append(
+                                        text_features_mlp.cpu().index_select(
+                                            0, torch.tensor(idx).long()
+                                        )
+                                    )
+                        #  print(f'eval step {i}') #  (yusong): for debug
 
                 # cumulative_loss += total_loss * batch_size
                 # num_samples += batch_size
-                if is_master(args) and (i % 100) == 0 and i != 0:
+                if is_master(args) and (i % 100) == 0:  # and i != 0:
                     logging.info(
                         f"Eval Epoch: {epoch} [{num_samples} / {samples_per_val}]"
                     )
-            val_metrics_s = {}
-            for n in eval_info.keys():
-                metrics_single_dataset = get_metrics(
-                    audio_features=torch.cat(eval_info[n]["all_audio_features"]),
-                    text_features=torch.cat(eval_info[n]["all_text_features"]),
-                    audio_features_mlp=torch.cat(
-                        eval_info[n]["all_audio_features_mlp"]
-                    ),
-                    text_features_mlp=torch.cat(eval_info[n]["all_text_features_mlp"]),
-                    logit_scale_a=logit_scale_a.cpu(),
-                    logit_scale_t=logit_scale_t.cpu(),
-                )
-                val_metrics_s[n] = {
-                    n + "/" + k: v for k, v in metrics_single_dataset.items()
-                }
-                metrics.update(val_metrics_s[n])
-                if "epoch" not in metrics.keys():
-                    metrics.update({"epoch": epoch})
-
-    if not metrics:
-        return metrics
-
-    logging.info(
-        f"Eval Epoch: {epoch} "
-        + "\n".join(
-            [
-                "\t".join([f"{k}: {round(v, 4):.4f}" for k, v in m.items()])
-                for m in val_metrics_s.values()
-            ]
+            if is_master(args):
+                val_metrics_per_dataset = {}
+                for n in eval_info.keys():
+                    if args.clap_mlploss:
+                        metrics_single_dataset = get_metrics(
+                            audio_features=torch.cat(eval_info[n]["all_audio_features"]),
+                            text_features=torch.cat(eval_info[n]["all_text_features"]),
+                            logit_scale_a=logit_scale_a.cpu(),
+                            audio_features_mlp=torch.cat(
+                                eval_info[n]["all_audio_features_mlp"]
+                            ),
+                            text_features_mlp=torch.cat(eval_info[n]["all_text_features_mlp"]),
+                            logit_scale_t=logit_scale_t.cpu(),
+                            mlp_loss=args.clap_mlploss
+                        )
+                    else:
+                        metrics_single_dataset = get_metrics(
+                            audio_features=torch.cat(eval_info[n]["all_audio_features"]),
+                            text_features=torch.cat(eval_info[n]["all_text_features"]),
+                            logit_scale_a=logit_scale_a.cpu(),
+                            mlp_loss=args.clap_mlploss
+                        )
+                    val_metrics_per_dataset[n] = {
+                        n + "/" + k: v for k, v in metrics_single_dataset.items()
+                    }
+                    metrics.update(val_metrics_per_dataset[n])
+                    if "epoch" not in metrics.keys():
+                        metrics.update({"epoch": epoch})
+    if is_master(args):
+        if not metrics:
+            return metrics
+
+        logging.info(
+            f"Eval Epoch: {epoch} "
+            + "\n".join(
+                [
+                    "\t".join([f"{k}: {round(v, 4):.4f}" for k, v in m.items()])
+                    for m in val_metrics_per_dataset.values()
+                ]
+            )
         )
-    )
 
-    if args.save_logs:
-        for name, val in metrics.items():
-            if tb_writer is not None:
-                tb_writer.add_scalar(f"val/{name}", val, epoch)
+        if args.save_logs:
+            for name, val in metrics.items():
+                if tb_writer is not None:
+                    tb_writer.add_scalar(f"val/{name}", val, epoch)
 
-        with open(os.path.join(args.checkpoint_path, "results.jsonl"), "a+") as f:
-            f.write(json.dumps(metrics))
-            f.write("\n")
+            with open(os.path.join(args.checkpoint_path, "results.jsonl"), "a+") as f:
+                f.write(json.dumps(metrics))
+                f.write("\n")
 
-    if args.wandb:
-        assert wandb is not None, "Please install wandb."
-        for name, val in metrics.items():
-            wandb.log({f"val/{name}": val, "epoch": epoch})
+        if args.wandb:
+            assert wandb is not None, "Please install wandb."
+            for name, val in metrics.items():
+                wandb.log({f"val/{name}": val, "epoch": epoch})
 
-    return metrics
+        return metrics
+    else:
+        return metrics
 
 
 def get_metrics(
-    audio_features,
-    text_features,
-    audio_features_mlp,
-    text_features_mlp,
-    logit_scale_a,
-    logit_scale_t,
+        audio_features,
+        text_features,
+        logit_scale_a,
+        audio_features_mlp=None,
+        text_features_mlp=None,
+        logit_scale_t=None,
+        mlp_loss=False
 ):
     metrics = {}
-    # Set up audio to text & text to audio similary matrice
-    a_logits_per_audio = (
-        (logit_scale_a * audio_features @ text_features_mlp.t()).detach().cpu()
-    )
-    a_logits_per_text = a_logits_per_audio.t().detach().cpu()
-    t_logits_per_audio = (
-        (logit_scale_t * audio_features_mlp @ text_features.t()).detach().cpu()
-    )
-    t_logits_per_text = t_logits_per_audio.t().detach().cpu()
-
-    labels = torch.arange(audio_features.shape[0]).long()
-    # Change the loss from two terms into four terms with 2x2 combined CE loss
-    total_loss = (
-        F.cross_entropy(a_logits_per_audio, labels)
-        + F.cross_entropy(a_logits_per_text, labels)
-        + F.cross_entropy(t_logits_per_audio, labels)
-        + F.cross_entropy(t_logits_per_text, labels)
-    ) / 4
-
-    metrics[f"cumulative_loss"] = total_loss.item()
-    metrics[f"num_samples"] = audio_features.shape[0]
-
-    logits = {
-        "audio_to_text": (a_logits_per_audio + t_logits_per_audio) / 2,
-        "text_to_audio": (a_logits_per_text + t_logits_per_text) / 2,
-    }
-    ground_truth = torch.arange(len(text_features)).view(-1, 1)
+    if mlp_loss:
+        # Set up audio to text & text to audio similary matrice
+        a_logits_per_audio = (
+            (logit_scale_a * audio_features @ text_features_mlp.t()).detach().cpu()
+        )
+        a_logits_per_text = a_logits_per_audio.t().detach().cpu()
+        t_logits_per_audio = (
+            (logit_scale_t * audio_features_mlp @ text_features.t()).detach().cpu()
+        )
+        t_logits_per_text = t_logits_per_audio.t().detach().cpu()
+
+        labels = torch.arange(audio_features.shape[0]).long()
+        # Change the loss from two terms into four terms with 2x2 combined CE loss
+        total_loss = (
+                             F.cross_entropy(a_logits_per_audio, labels)
+                             + F.cross_entropy(a_logits_per_text, labels)
+                             + F.cross_entropy(t_logits_per_audio, labels)
+                             + F.cross_entropy(t_logits_per_text, labels)
+                     ) / 4
+
+        metrics[f"cumulative_loss"] = total_loss.item()
+        metrics[f"num_samples"] = audio_features.shape[0]
+
+        logits = {
+            "audio_to_text": (a_logits_per_audio + t_logits_per_audio) / 2,
+            "text_to_audio": (a_logits_per_text + t_logits_per_text) / 2,
+        }
+        ground_truth = torch.arange(len(text_features)).view(-1, 1)
+
+    else:
+        # print("text_features", text_features)
+        # print("text_features.shape", text_features.shape)
+        logits_per_audio = (logit_scale_a * audio_features @ text_features.t()).detach().cpu()
+        logits_per_text = logits_per_audio.t().detach().cpu()
+
+        labels = torch.arange(audio_features.shape[0]).long()
+        # Change the loss from two terms into four terms with 2x2 combined CE loss
+        total_loss = (
+                             F.cross_entropy(logits_per_audio, labels)
+                             + F.cross_entropy(logits_per_text, labels)
+                     ) / 2
+
+        metrics[f"cumulative_loss"] = total_loss.item()
+        metrics[f"num_samples"] = audio_features.shape[0]
+
+        logits = {"audio_to_text": logits_per_audio, "text_to_audio": logits_per_text}
+
+        ground_truth = torch.arange(len(text_features)).view(-1, 1)
 
     for name, logit in logits.items():
         ranking = torch.argsort(logit, descending=True)
-        preds = torch.where(ranking == ground_truth)[
-            1
-        ]  # (yusong) this line is slow because it uses single thread
+        preds = torch.where(ranking == ground_truth)[1]  # (yusong) this line is slow because it uses single thread
         preds = preds.detach().cpu().numpy()
         metrics[f"{name}_mean_rank"] = preds.mean() + 1
         metrics[f"{name}_median_rank"] = np.floor(np.median(preds)) + 1
@@ -424,3 +573,208 @@ def get_metrics(
         metrics[f"{name}_mAP@10"] = np.mean(np.where(preds < 10, 1 / (preds + 1), 0.0))
 
     return metrics
+
+
+def evaluate_clotho_audiocaps(
+        model, data, epoch, args, autocast, device, tb_writer=None
+):
+    """
+    Adapted from https://github.com/XinhaoMei/audio-text_retrieval/blob/main/tools/utils.py.
+    1. for text-to-audio retrieval, do 5 times and average the results
+    2. for R@1, R@5, R@10 in audio-to-text retrieval, take the best rank among 5 text
+    3. for map@10 in audio-to-text retrieval:
+        3.1: sort the rank of 5 text
+        3.2: exclude the rank >=10 (0-index)
+        3.3: compute the map regarding the remaining ranks: np.mean(np.arange(1, len(ranks)+1) / ranks).
+        (3.3) That is, take the top ranks of 5 text that is < 10, and assign the descending number as ground truth.
+        (3.3) E.g.: the ground truth of first rank of the 5 text should be 1, the second rank should be 2, etc.
+    """
+    # TODO: (yusong) only support single GPU evaluation and only support non-mlp case for now.
+    dataloader = data["val"].dataloader
+    with torch.no_grad():
+        eval_info = {}
+        for i, batch in enumerate(dataloader):
+            audios = batch  # contains mel_spec, wavform, and longer list
+
+            # each item in the list has 5 texts
+            if args.tmodel == "transformer":
+                from open_clip import tokenize
+                texts = [tokenize(t) for t in batch['full_text']]
+                texts = torch.cat(texts)
+            else:
+                from .data import tokenizer
+                texts = [tokenizer(t) for t in batch['full_text']]  # 5 texts for each audio
+                texts = {k: torch.cat([t[k] for t in texts]) for k in texts[0].keys()}  # 5 x batch
+
+            # audios = audios.to(device=device, non_blocking=True)
+
+            all_names = list(set(["-".join(b.split("/")[-3:-1]) for b in batch['__url__']]))
+            for name in all_names:
+                if name not in eval_info.keys():
+                    # we will not use mlp outputs even if args.clap_mlploss=True
+                    eval_info[name] = {
+                        "cumulative_loss": 0.0,
+                        "num_samples": 0,
+                        "all_audio_features": [],
+                        "all_text_features": []
+                    }
+            with autocast():
+                audio_features = model(audios, None, device)
+                text_features = model(None, texts, device)
+                audio_features = F.normalize(audio_features, dim=-1)
+                text_features = F.normalize(text_features, dim=-1)
+
+                all_names = list(set(["-".join(b.split("/")[-3:-1]) for b in batch['__url__']]))
+                for n in all_names:
+                    idx = np.where(
+                        np.array(
+                            ["-".join(b.split("/")[-3:-1]) for b in batch['__url__']]
+                        )
+                        == n
+                    )[0]
+                    eval_info[n]["all_audio_features"].append(
+                        audio_features.cpu().index_select(
+                            0, torch.tensor(idx).long()
+                        )
+                    )
+                    # (yusong) please double-check. This is for selecting 5 text features at once.
+                    # because idx is a list of indices in size of num_samples,
+                    # and text_features is a tensor of size (5*num_samples, dim)
+                    # so we need to select 5 consecutive indices at once for a single index in idx.
+                    eval_info[n]["all_text_features"].append(
+                        text_features.cpu().reshape([-1, 5, text_features.shape[1]]).index_select(
+                            0, torch.tensor(idx).long()
+                        ).reshape([-1, text_features.shape[1]])
+                    )
+
+        val_metrics_all = {}
+
+        for n in eval_info.keys():
+            logit_scale_a, logit_scale_t = model(None, None, device)
+            logit_scale_a = logit_scale_a.cpu()
+
+            audio_features = torch.cat(eval_info[n]["all_audio_features"], dim=0)
+            text_features = torch.cat(eval_info[n]["all_text_features"], dim=0)
+
+            logits_per_audio = (logit_scale_a * audio_features @ text_features.t()).detach().cpu()
+            logits_per_text = logits_per_audio.t().detach().cpu()
+
+            # logits_per_audio shape: [num_samples, num_samples*5]
+            # logits_per_text shape: [num_samples*5, num_samples]
+
+            logging.info(f"dataset {n}, logits_per_audio shape: {logits_per_audio.shape}, "
+                         f"logits_per_text shape: {logits_per_text.shape}")
+
+            metrics = {}
+            num_samples = audio_features.shape[0]
+            metrics[f"num_samples"] = num_samples
+
+            # (yusong) the following code is very important, please double-check:
+            # logits_per_audio.reshape(num_samples, num_samples, 5)[:, :, d]
+            # logits_per_text.reshape(num_samples, 5, num_samples)[:, d, :]
+            # Those two are retrieving one of the 5 text for each audio.
+            labels = torch.arange(audio_features.shape[0]).long()
+            audio_to_text_loss = [
+                F.cross_entropy(
+                    logits_per_audio.reshape(num_samples, num_samples, 5)[:, :, d], labels) for d in range(5)
+            ]
+            text_to_audio_loss = [
+                F.cross_entropy(
+                    logits_per_text.reshape(num_samples, 5, num_samples)[:, d, :], labels) for d in range(5)
+            ]
+            total_loss = (
+                                 np.mean(audio_to_text_loss) + np.mean(text_to_audio_loss)
+                         ) / 2
+
+            metrics[f"cumulative_loss"] = total_loss.item()
+
+            # text to audio: do 5 times
+            pred_text = []
+            for d in range(5):
+                logit = logits_per_text.reshape(num_samples, 5, num_samples)[:, d, :]
+                ground_truth = torch.arange(len(logit)).view(-1, 1)
+                ranking = torch.argsort(logit, descending=True)  # [num_samples, num_samples]
+                preds = torch.where(ranking == ground_truth)[1]
+                pred_text.append(preds.detach().cpu().numpy())
+            pred_text_concat = np.concatenate(pred_text, axis=0)  # [5*num_samples]
+            metrics[f"text_to_audio_mean_rank"] = pred_text_concat.mean() + 1
+            metrics[f"text_to_audio_median_rank"] = np.floor(np.median(pred_text_concat)) + 1
+            for k in [1, 5, 10]:
+                metrics[f"text_to_audio_R@{k}"] = np.mean(pred_text_concat < k)
+            # map@10
+            metrics[f"text_to_audio_mAP@10"] = np.mean(np.where(pred_text_concat < 10, 1 / (pred_text_concat + 1), 0.0))
+
+            # audio to text: take the best result
+            # for audio to text map 10, sort and assign descending ground truth.
+            # see https://github.com/XinhaoMei/audio-text_retrieval/blob/main/tools/utils.py#L103
+            # map@10
+            map_all = []
+            pred_audio_all = []
+            for d in range(num_samples):
+                # logits_per_audio: [num_samples, num_samples*5]
+                logit_single = logits_per_audio[d, :]  # [5*num_samples]
+                # Ground-truth index: [d*5, d*5+1, d*5+2, d*5+3, d*5+4]
+                ranking = torch.argsort(logit_single, descending=True)  # [5*num_samples]
+                # ranking: the index of first match, second match, ...
+                ground_truth = torch.arange(d * 5, d * 5 + 5)[None]
+                all_pred = torch.where(torch.stack([ranking] * 5) == ground_truth.view(-1, 1))[1]
+                min_pred = torch.min(all_pred)
+                pred_audio_all.append(min_pred.detach().cpu().numpy())
+                all_pred_filter = all_pred[all_pred < 10].detach().cpu().numpy()
+                # /5 because we have 5 text, so it means for the text rank >=10 we count as 0.
+                map_single = np.sum((np.arange(1, len(all_pred_filter) + 1) / (all_pred_filter + 1))) / 5
+                map_all.append(map_single)
+            metrics[f"audio_to_text_mAP@10"] = np.mean(map_all)
+            for k in [1, 5, 10]:
+                metrics[f"audio_to_text_R@{k}"] = np.mean(np.array(pred_audio_all) < k)
+
+            val_metrics_all[n] = {
+                n + "/" + k: v for k, v in metrics.items()
+            }
+    return val_metrics_all
+
+
+def calculate_selection_performance_clotho_audiocaps(val_metrics_per_dataset):
+    """
+    Calculate performance for Clotho+AudioCaps for model selection.
+    """
+    selection_performance_all = []
+    for n in val_metrics_per_dataset.keys():
+        selection_performance = (val_metrics_per_dataset[n][f"{n}/audio_to_text_mAP@10"] +
+                                 val_metrics_per_dataset[n][f"{n}/text_to_audio_mAP@10"]) / 2
+        selection_performance_all.append(selection_performance)
+    return np.mean(selection_performance_all)
+
+
+def select_top_metric_clotho_audiocaps(metrics, val_metrics_per_dataset, args):
+    # val_metrics_per_dataset: dict, key: dataset name, value: dict, key: metric name, value: metric value
+    # metrics: dict, key: metric name, value: metric value
+    # Hack: use args to save the top performance
+    if not hasattr(args, "top_selection_performance"):
+        selection_performance = calculate_selection_performance_clotho_audiocaps(val_metrics_per_dataset)
+        # TODO: write the if and else together
+        metric_update = {}
+        for n in val_metrics_per_dataset.keys():
+            for k in val_metrics_per_dataset[n].keys():
+                metric_update[k.split('/')[0] + '-top' + '/' + k.split('/')[1]] = val_metrics_per_dataset[n][k]
+        metric_update['top_selection_performance'] = selection_performance
+        metric_update['top-selection-epoch'] = metrics['epoch']
+        metrics.update(metric_update)
+        args.top_metric = metric_update
+        args.top_selection_performance = selection_performance
+    else:
+        selection_performance_new = calculate_selection_performance_clotho_audiocaps(val_metrics_per_dataset)
+        selection_performance_old = args.top_selection_performance
+        if selection_performance_new > selection_performance_old:
+            metric_update = {}
+            for n in val_metrics_per_dataset.keys():
+                for k in val_metrics_per_dataset[n].keys():
+                    metric_update[k.split('/')[0] + '-top' + '/' + k.split('/')[1]] = val_metrics_per_dataset[n][k]
+            metric_update['top_selection_performance'] = selection_performance_new
+            metric_update['top-selection-epoch'] = metrics['epoch']
+            metrics.update(metric_update)
+            args.top_metric = metric_update
+            args.top_selection_performance = selection_performance_new
+        else:
+            metrics.update(args.top_metric)
+    return metrics


Test patch:
diff --git a/src/test_tars.py b/src/test_tars.py
new file mode 100644
index 0000000..f1bef55
--- /dev/null
+++ b/src/test_tars.py
@@ -0,0 +1,120 @@
+import webdataset as wds
+import soundfile as sf
+import io
+import os
+import random
+import copy
+from tqdm import tqdm
+import shutil
+import argparse
+import traceback
+import logging
+import json
+from open_clip import tokenize
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--tar-path",
+        type=str,
+        default=None,
+        help="Path to the tars",
+    )
+    parser.add_argument(
+        "--start",
+        type=int,
+        default=0,
+        help="start from tar-path + start",
+    )
+    parser.add_argument(
+        "--end",
+        type=int,
+        default=99999,
+        help="end with tar-path + end",
+    )
+    parser.add_argument(
+        "--exclude",
+        nargs='+',
+        default=None,
+        help="exclude tar-path + exclude",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=1,
+    )
+    parser.add_argument(
+        "--order",
+        default=False,
+        action='store_true',
+        help="if keep the search order accendingly",
+    )
+    args = parser.parse_args()
+    return args
+
+def log_and_continue(exn):
+    """Call in an exception handler to ignore any exception, isssue a warning, and continue."""
+    logging.warning(f"Handling webdataset error ({repr(exn)}). Ignoring.")
+    return True
+
+def preprocess(
+    sample,
+):
+    """
+    Preprocess a single sample for wdsdataloader.
+    """
+    audio_ext = "flac"
+    text_ext = "json"
+    audio_data, orig_sr = sf.read(io.BytesIO(sample[audio_ext]))
+    json_dict_raw = json.loads(sample[text_ext].decode("utf-8"))
+    sample["waveform"] = audio_data
+    texts = json_dict_raw["text"]
+    if isinstance(texts, list) and isinstance(texts[0], str) and len(texts) > 1:
+        texts = random.choice(texts)
+    sample["raw_text"] = texts
+    sample["text"] = tokenize(texts)
+    return sample
+
+if __name__ == "__main__":
+    args = parse_args()
+    tar_path = args.tar_path
+    idx_list = list(range(args.start, args.end))
+    if args.exclude != None:
+        for x in args.exclude:
+            idx_list.remove(x)
+    if not args.order:
+        random.shuffle(idx_list)
+    if "aws" in tar_path:
+        args.local = False
+    if args.local:
+        input_shards = [os.path.join(args.tar_path, str(i)+".tar") for i in idx_list]
+    else:
+        input_shards = [os.path.join(args.tar_path, str(i)+".tar -") for i in idx_list]
+    pipeline = [wds.SimpleShardList(input_shards)]
+    pipeline.extend(
+        [
+            wds.split_by_node,
+            wds.split_by_worker,
+            wds.tarfile_to_samples(handler=log_and_continue),
+            wds.map(preprocess),
+            wds.to_tuple("__url__", "__key__", "waveform"),
+            wds.batched(1),
+        ]
+    )
+    dataset = wds.DataPipeline(*pipeline)
+    dataloader = wds.WebLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=0)
+    old_k = 0
+    old_batch = None
+    try:
+        for k, batch in tqdm(enumerate(dataloader)):
+            print("k:", k)
+            print("batch:", batch)
+            old_k = k
+            old_batch = copy.deepcopy(batch)
+    except:
+        with open("check_tar_log.txt","a") as file:
+            traceback.print_exc(file = file)
+        print("old_k:", old_k)
+        print("old_batch:", old_batch)
+        pass
diff --git a/src/training/test_infer.py b/src/training/test_infer.py
index 924c417..2c6b7cc 100644
--- a/src/training/test_infer.py
+++ b/src/training/test_infer.py
@@ -127,7 +127,7 @@ def random_seed(seed=42, rank=0):
 def main():
     args = parse_args()
     # sanitize model name for filesystem / uri use, easier if we don't use / in name as a rule?
-    args.model = args.model.replace("/", "-")
+    args.amodel = args.amodel.replace("/", "-")
     # download sizes.json file
     
     random.seed(args.seed)
@@ -141,7 +141,7 @@ def main():
         args.name = "-".join(
             [
                 datetime.now().strftime("%Y_%m_%d-%H_%M_%S"),
-                f"model_{args.model}",
+                f"model_{args.amodel}",
                 f"lr_{args.lr}",
                 f"b_{args.batch_size}",
                 f"j_{args.workers}",
@@ -193,7 +193,7 @@ def main():
         logging.info(f"Running with a single process. Device {args.device}.")
 
     model, model_cfg = create_model(
-        args.model,
+        args.amodel,
         args.pretrained,
         precision=args.precision,
         device=device,
@@ -249,7 +249,7 @@ def main():
     ]
 
     # construct data
-    output_dict = model.audio_infer(torch.rand(960000).cuda(), hopsize=120000, key ="fine_grained_embedding")
+    output_dict = model.audio_infer(torch.rand(960000).cuda(), hopsize=120000, key ="fine_grained_embedding", device=device)
     print(output_dict["fine_grained_embedding"].size())
 
     # inference