Changes
1 changed files (+653/-0)
-
src/main.rs.best (new)
-
@@ -0,0 +1,653 @@use rand::Rng; use rand::distr::Distribution; use rand::distr::weighted::WeightedIndex; use rand_xoshiro::Xoshiro256PlusPlus; use rand_xoshiro::rand_core::SeedableRng; use std::cmp::Ordering; use std::collections::{HashSet, VecDeque}; use std::env; use std::fs::File; use std::io::{BufReader, prelude::*}; use std::thread; use tch::nn::{Module, OptimizerConfig}; use tch::{Device, IndexOp, Kind, NewAxis, Reduction, Tensor, nn, no_grad}; const LR: f64 = 5e-4; const WD: f64 = 1e-4; const BS: usize = 4096; const EPOCHS: usize = 20; const SEQ_LEN: usize = 256; const EMBED: i64 = 12; const HIDDEN: i64 = 512; const KERNEL: i64 = 16; const WINDOW: usize = 30; const ITERS: usize = 15000; // Iters for length 100 const THREADS: usize = 4; const VOCAB: usize = 28; // Sorted from lowest to highest freq const VOCAB_ORDER: [u8; VOCAB] = [ 25, 16, 9, 23, 10, 21, 27, 1, 15, 6, 24, 5, 2, 22, 12, 20, 11, 3, 17, 18, 7, 8, 13, 14, 0, 19, 4, 26, ]; #[rustfmt::skip] const TRANS: [[f32; VOCAB]; VOCAB] = [ [-7.591_938_5, -2.892_999_6, -2.175_499, -3.498_550_7, -2.996_972_6, -2.549_879_6, -2.822_673_3, -1.661_220_1, -4.226_744_7, -1.847_851_5, -5.350_03, -2.655_289_2, -2.019_139, -4.166_14, -5.359_595_3, -2.123_517_8, -11.513_105, -2.907_192, -3.031_055_5, -3.436_183, -3.792_863_6, -2.951_992_8, -1.479_117_6, -2.626_861_3, -5.394_068, -0.961_980_4, -2.162_345, -20.723_267], [-3.829_590_8, -5.279_662, -9.647_869, -7.151_901_2, -6.973_281, -8.775_774, -6.730_87, -6.306_915_3, -5.006_757_7, -11.513_105, -6.914_897_4, -7.158_250_3, -4.238_712, -7.696_574, -4.852_602, -6.468_646, -11.513_105, -6.261_517_5, -6.552_708_6, -8.077_111, -4.138_724_3, -11.513_105, -4.975_327_5, -8.290_131, -5.404_796_6, -11.513_105, -3.067_354, -20.723_267], [-3.567_029, -10.307_745, -3.879_610_8, -8.355_694, -3.945_243_6, -7.539_704, -8.209_732, -9.787_797, -3.231_729, -11.513_105, -7.435_984, -6.633_628_4, -8.991_505, -3.292_566_8, -4.787_024, -10.329_214, -11.513_105, -4.649_024, -4.578_408_2, -5.882_728_6, -2.919_008_5, -11.513_105, -7.183_579, -1.746_624_2, -7.366_245_7, -7.672_502_5, -3.292_724_4, -20.723_267], [-2.840_364_5, -7.370_669, -9.647_869, -4.581_870_6, -2.539_729, -7.668_403_6, -7.623_522, -8.273_909, -3.187_738, -11.513_105, -8.247_089, -2.392_907_1, -8.366_328, -1.767_054_9, -4.048_143, -9.203_718, -11.513_105, -3.572_408_7, -7.468_721, -9.571_307, -4.192_383_3, -11.513_105, -6.453_733, -8.944_214, -6.014_779, -11.513_105, -3.493_481, -20.723_267], [-8.789_968_5, -0.882_509_6, -1.761_314_5, -2.229_609, -3.460_188_2, -2.379_223_8, -2.100_699, -0.839_311_1, -3.267_797_2, -1.279_895, -1.137_539_4, -1.836_477_3, -1.454_008, -2.393_423_3, -5.954_031_5, -1.575_469_5, -11.513_105, -1.506_992_6, -2.176_778_3, -2.483_203, -3.771_802, -0.265_394_7, -1.942_219, -2.516_527_7, -3.196_823, -1.083_491_7, -3.787_814_6, -20.723_267], [-4.656_895_6, -10.307_745, -11.513_105, -6.434_101, -4.661_793, -3.189_781, -8.278_472, -7.476_047, -3.797_279_8, -11.513_105, -4.357_318, -3.576_648_7, -5.378_657, -5.115_966_3, -2.389_812_7, -8.831_921, -11.513_105, -4.785_914_4, -5.850_82, -6.378_231_5, -5.453_339, -11.513_105, -5.192_757, -5.090_478, -6.588_798_5, -11.513_105, -3.367_658, -20.723_267], [-3.935_778_6, -11.513_105, -11.513_105, -5.198_305, -5.182_173_3, -9.076_6, -4.308_548, -9.307_589, -3.593_207_8, -11.513_105, -7.141_876_7, -6.870_805_7, -10.692_42, -2.028_623, -5.771_454, -9.323_16, -11.513_105, -5.052_397_3, -6.821_178, -9.731_103, -2.909_238_3, -11.513_105, -9.689_901, -8.944_214, -8.006_65, -11.513_105, -4.130_123, -20.723_267], [-6.750_226, -8.066_504, -1.719_597, -6.110_409_7, -6.189_956_7, -7.899_103_6, -1.912_184_1, -9.160_122, -9.226_412, -11.513_105, -5.135_840_4, -8.518_253, -8.545_525, -6.787_862, -5.718_264_6, -4.609_116, -11.513_105, -5.813_02, -2.525_067_6, -1.224_304_8, -10.253_361, -11.513_105, -1.766_828_5, -4.545_65, -7.245_962_6, -7.672_502_5, -2.483_950_6, -20.723_267], [-3.126_791_7, -3.799_575, -3.186_328_2, -2.768_115_3, -4.389_353, -2.612_637, -3.096_265_6, -2.015_972_1, -8.032_84, -3.987_442, -1.910_051_2, -2.178_937, -2.289_748, -3.666_899_2, -4.623_344, -2.913_505_8, -11.513_105, -2.678_056_2, -2.861_862, -2.780_634, -3.503_753, -2.036_225_6, -1.760_814_5, -2.264_577_6, -4.232_542, -2.710_460_4, -2.579_534_8, -20.723_267], [-9.585_668, -4.257_89, -11.513_105, -8.256_346, -7.803_908_3, -11.513_105, -10.480_922, -11.057_719, -11.513_105, -11.513_105, -9.365_77, -9.140_743, -10.692_42, -6.505_830_3, -9.031_633, -10.755_326, -11.513_105, -9.805_901, -10.500_641, -10.708_67, -10.696_741, -11.513_105, -10.603_504, -11.513_105, -10.070_196, -11.513_105, -5.602_34, -20.723_267], [-4.254_719_7, -11.513_105, -3.200_003, -8.096_455, -6.754_756_5, -10.923_909, -9.805_095, -10.508_677_5, -4.890_956, -11.513_105, -8.732_795, -4.552_042, -11.020_834, -4.463_423_3, -4.283_471_6, -8.257_408, -11.513_105, -5.009_569_6, -5.384_791_4, -11.325_874, -8.812_76, -11.513_105, -7.243_991, -11.513_105, -10.551_27, -11.513_105, -5.107_608, -20.723_267], [-2.719_842, -2.171_247_2, -3.465_415_2, -4.541_519, -3.172_815_8, -4.154_808_5, -3.646_469_4, -6.766_412, -3.120_190_6, -11.513_105, -4.313_071_3, -1.951_761_5, -5.967_743, -4.602_898_6, -3.833_565, -2.322_928_4, -11.513_105, -4.435_82, -5.262_661, -4.106_762_4, -2.086_468_5, -11.513_105, -5.110_801_7, -6.713_080_4, -5.764_374_7, -3.582_983_7, -3.718_748_6, -20.723_267], [-3.592_866_7, -6.562_377_5, -10.936_696, -5.287_308, -3.877_893, -9.623_346, -6.248_005_4, -6.303_758, -3.124_795_2, -11.513_105, -7.080_122_5, -5.256_596_6, -3.885_822_5, -7.657_765_4, -2.880_621, -7.727_230_5, -11.513_105, -4.395_987, -5.172_2, -6.633_761, -4.039_342, -11.513_105, -8.661_436, -11.513_105, -5.230_367_7, -11.513_105, -2.871_839_8, -20.723_267], [-1.597_543_4, -9.295_766, -10.097_605, -5.292_990_7, -2.487_039_3, -9.313_064, -4.270_337_6, -6.744_131_6, -1.308_721_7, -11.513_105, -1.928_257_2, -5.963_909, -6.134_99, -4.110_112_7, -1.975_486_4, -9.323_16, -11.513_105, -3.969_218_3, -6.406_316_3, -6.633_761, -2.287_234_3, -10.490_972, -3.211_650_1, -11.513_105, -6.668_236, -11.513_105, -3.653_006, -20.723_267], [-7.741_305_4, -2.454_400_5, -1.525_927, -2.863_657_2, -6.311_463, -1.725_130_9, -2.603_309_2, -2.540_201_4, -3.049_867, -1.419_446_5, -6.538_593, -2.566_781_3, -2.362_048_6, -2.469_979_5, -3.197_990_2, -2.080_492, -11.513_105, -2.697_022_7, -2.865_752, -2.179_993, -6.900_775, -3.221_613_4, -2.312_064_2, -4.879_478_5, -1.818_898_6, -4.474_464, -2.843_442_7, -20.723_267], [-4.030_286_3, -11.513_105, -10.573_381, -9.024_024, -4.715_614_3, -7.836_251, -7.501_858, -9.787_797, -5.341_391, -11.513_105, -8.247_089, -5.940_762_5, -3.205_545_4, -7.367_222, -4.124_150_8, -2.470_337_9, -11.513_105, -5.421_519_3, -3.863_525_9, -8.861_415, -3.146_821_3, -11.513_105, -8.619_498, -1.339_567_1, -7.223_549_4, -11.513_105, -3.670_306_2, -20.723_267], [-10.869_481, -11.513_105, -5.161_777_5, -10.273_109, -6.273_897, -10.923_909, -10.869_361, -10.620_307, -8.873_802, -11.513_105, -11.513_105, -10.857_676, -11.513_105, -6.483_735, -9.112_832, -10.755_326, -11.513_105, -8.188_78, -7.532_092, -10.912_28, -10.088_636, -11.513_105, -11.513_105, -4.863_750_5, -10.551_27, -11.513_105, -6.010_124_7, -20.723_267], [-2.380_518_7, -2.971_973_2, -3.193_902_3, -4.207_732, -1.864_907_7, -2.681_471_8, -2.695_279_8, -4.959_060_7, -3.280_368_3, -11.513_105, -8.155_254, -6.200_581_6, -2.820_928_8, -6.244_217, -2.220_730_8, -1.956_054, -11.513_105, -3.837_347, -7.571_002_5, -3.770_41, -1.872_708_4, -6.901_649, -4.567_485, -11.513_105, -7.268_889_4, -11.513_105, -3.974_993, -20.723_267], [-2.199_485_3, -4.014_287, -7.049_959, -3.797_005_2, -2.864_42, -7.275_875_6, -3.802_058_2, -6.339_048_4, -2.140_158_4, -11.513_105, -3.389_352_3, -4.379_51, -3.850_523_7, -3.391_756_5, -3.693_605_7, -3.922_179, -11.513_105, -2.862_128_3, -2.919_803_9, -4.157_688, -2.148_169, -11.513_105, -4.797_732_4, -5.466_213, -3.208_151, -11.513_105, -2.598_197_7, -20.723_267], [-2.041_496_8, -4.551_41, -2.453_904_6, -7.047_818, -3.565_677, -3.377_406_8, -5.630_275, -3.469_57, -1.993_827_8, -11.513_105, -7.791_498_7, -4.063_527, -7.335_05, -2.424_767_5, -2.803_703_8, -3.327_58, -11.513_105, -3.117_431_9, -2.197_329_8, -3.673_175_6, -2.026_581_5, -11.513_105, -7.587_344, -1.843_796_1, -4.424_683_6, -6.990_154, -2.010_327_6, -20.723_267], [-4.493_191_7, -2.045_081_9, -3.441_515_4, -4.625_169_8, -7.377_274, -3.366_884_7, -4.019_038_7, -4.824_732_3, -7.603_225_7, -1.204_362, -8.348_22, -4.263_171, -3.008_42, -5.452_8, -1.893_440_6, -3.538_099, -0.001_714_452_5, -4.397_697, -3.250_448_5, -4.170_668_6, -10.450_676, -6.800_828_5, -8.467_393, -4.592_125, -10.070_196, -5.612_035_3, -4.575_990_7, -20.723_267], [-3.431_514_7, -8.532_299, -11.513_105, -5.830_799_6, -3.924_140_2, -10.555_642, -11.513_105, -10.889_751, -3.970_075_8, -11.513_105, -10.254_454, -5.386_309_6, -10.247_817, -5.416_301_3, -4.264_887, -11.513_105, -11.513_105, -4.975_447_7, -9.886_048, -10.912_28, -7.955_444_3, -11.513_105, -10.958_263, -5.653_776, -6.996_506, -7.672_502_5, -4.962_283_6, -20.723_267], [-4.383_549, -9.175_94, -11.513_105, -6.060_079_6, -5.145_225, -8.170_727, -7.974_089_6, -6.377_324_6, -10.387_997, -11.513_105, -5.921_405_3, -4.953_966_6, -9.043_949, -6.696_472, -3.019_770_4, -7.606_258_4, -11.513_105, -5.849_935, -5.396_963_6, -4.984_56, -9.293_391, -11.513_105, -8.849_356, -7.218_435, -6.520_682_3, -4.082_706_5, -2.607_651, -20.723_267], [-7.124_895, -11.513_105, -11.513_105, -11.513_105, -4.681_599, -11.513_105, -11.513_105, -11.513_105, -6.373_989, -11.513_105, -11.513_105, -11.513_105, -11.513_105, -6.898_58, -7.228_384, -11.513_105, -11.513_105, -11.513_105, -11.513_105, -11.513_105, -7.596_529_5, -11.513_105, -11.513_105, -5.466_213, -10.920_747, -11.513_105, -9.556_946, -20.723_267], [-3.479_500_8, -2.599_211_7, -4.462_72, -4.215_739_3, -4.337_405_7, -6.045_270_4, -6.257_39, -5.687_123, -11.281_939, -11.513_105, -5.127_308, -2.319_987_3, -2.680_487_9, -4.002_873_4, -5.709_786, -4.407_800_7, -8.677_725, -3.129_120_3, -5.966_902, -4.185_307, -7.634_954, -4.926_794_5, -7.647_650_2, -6.452_617, -8.497_671, -2.881_854_5, -4.042_910_6, -20.723_267], [-7.170_823_6, -11.513_105, -11.513_105, -11.513_105, -9.303_842, -11.513_105, -10.201_862, -11.513_105, -5.639_653_7, -11.513_105, -11.513_105, -11.513_105, -11.513_105, -9.832_733, -8.523_063, -11.513_105, -11.513_105, -10.723_01, -11.513_105, -9.001_49, -7.790_009_5, -11.513_105, -10.958_263, -11.513_105, -10.920_747, -2.501_138_4, -10.541_421, -20.723_267], [-2.715_362_5, -5.991_672_5, -5.186_562_5, -0.495_729_1, -1.120_387_1, -0.994_527_76, -1.014_298_6, -2.282_205, -2.596_138_5, -11.513_105, -1.262_556_9, -2.028_100_3, -1.974_505_3, -1.473_895_5, -1.899_170_9, -2.886_320_6, -8.677_725, -1.367_780_4, -1.043_720_5, -1.285_834_2, -2.472_834_8, -7.485_031, -2.315_513_6, -2.507_613_4, -0.409_93, -4.361_229, -20.723_267, 1e-09], [-7.057_105, -7.864_855_3, -6.832_732_7, -3.767_181_4, -4.200_749, -4.790_685, -4.105_086, -5.590_689, -7.180_27, -5.557_085, -4.052_763, -4.921_688_6, -4.251_131, -4.469_908_7, -6.025_704, -5.421_922, -6.783_435, -3.429_898, -3.662_474, -4.345_318, -5.380_872_7, -8.071_246, -5.128_164_3, -4.732_333_7, -3.275_528_7, -11.513_105, -20.723_267, -20.723_267], ]; /// Conv2D with middle element masked #[derive(Debug)] struct MaskedConv { pub convl: nn::Conv2D, pub convr: nn::Conv2D, } impl MaskedConv { fn new( vs: &nn::Path, in_dim: i64, out_dim: i64, kernel_dim: [i64; 2], config: nn::ConvConfigND<[i64; 2]>, ) -> MaskedConv { MaskedConv { convl: nn::conv(vs, in_dim, out_dim, kernel_dim, config), convr: nn::conv(vs, in_dim, out_dim, kernel_dim, config), } } } impl Module for MaskedConv { fn forward(&self, xs: &Tensor) -> Tensor { // tch-rs doesn't like negative indices in slices (self.convl.forward(&xs.i((.., NewAxis, KERNEL + 1.., ..))) + self .convr .forward(&xs.i((.., NewAxis, ..xs.size()[1] - KERNEL - 1, ..)))) .squeeze() .transpose(-2, -1) } } fn net(vs: &nn::Path) -> impl Module + use<> { nn::seq() .add(nn::embedding(vs, VOCAB as i64, EMBED, Default::default())) .add(MaskedConv::new( vs, 1, HIDDEN, [KERNEL, EMBED], Default::default(), )) .add_fn(|xs| xs.relu()) .add(nn::linear(vs, HIDDEN, HIDDEN, Default::default())) .add_fn(|xs| xs.relu()) .add(nn::linear(vs, HIDDEN, HIDDEN, Default::default())) .add_fn(|xs| xs.relu()) .add(nn::linear(vs, HIDDEN, VOCAB as i64, Default::default())) .add_fn(|xs| xs.log_softmax(-1, Kind::Float)) } fn probs(net: &impl Module, s: &[u8]) -> (Tensor, f64) { let data: Vec<_> = s.iter().map(|x| *x as i64).collect(); // We only do inference on CPU so no need to move to device let xs = Tensor::from_slice2(&[data]); // Remove batch dim let ys = net.forward(&xs).squeeze(); let loss = ys .nll_loss(&xs.i((0, KERNEL..s.len() as i64 - KERNEL))) .double_value(&[]); (ys, loss) } /// Pad both sides with spaces fn probs_padded(net: &impl Module, s: &[u8]) -> (Tensor, f64) { probs( net, &[&[26_u8; KERNEL as usize], s, &[26_u8; KERNEL as usize]].concat(), ) } fn char_to_label(c: char) -> u8 { if c.is_ascii_lowercase() { c as u8 - b'a' } else if c == ' ' { 26 } else { 27 } } fn label_to_char(l: u8) -> char { match l.cmp(&26) { Ordering::Less => (b'a' + l) as char, Ordering::Equal => ' ', Ordering::Greater => '.', } } fn to_string(s: &[u8]) -> String { s.iter().map(|l| label_to_char(*l)).collect() } fn permute(s: &[u8], p: &[u8; VOCAB]) -> Vec<u8> { s.iter().map(|l| p[*l as usize]).collect() } fn logprob( s: &[u8], p: &[u8; VOCAB], cnts: &[[i32; VOCAB]; VOCAB], grams: &[[[f32; VOCAB]; VOCAB]; VOCAB], ) -> f32 { let mut lp = 0.; if s.len() < 300 { // Use trigrams for short and hard sequences for i in 1..s.len() { lp += TRANS[p[s[i] as usize] as usize][p[s[i - 1] as usize] as usize]; } lp += TRANS[p[s[1] as usize] as usize][p[s[0] as usize] as usize]; for i in 2..s.len() { lp += grams[p[s[i - 2] as usize] as usize][p[s[i - 1] as usize] as usize] [p[s[i] as usize] as usize] / 2.; } } else { for i in 0..VOCAB { for j in 0..VOCAB { lp += cnts[i][j] as f32 * TRANS[p[j] as usize][p[i] as usize]; } } } // unsafe { // if s.len() < 300 { // // Use trigrams for short and hard sequences // for i in 1..s.len() { // lp += TRANS // .get_unchecked(*p.get_unchecked(s[i] as usize) as usize) // .get_unchecked(*p.get_unchecked(s[i - 1] as usize) as usize); // } // lp += TRANS[p[s[1] as usize] as usize][p[s[0] as usize] as usize]; // for i in 2..s.len() { // lp += grams // .get_unchecked(*p.get_unchecked(s[i - 2] as usize) as usize) // .get_unchecked(*p.get_unchecked(s[i - 1] as usize) as usize) // .get_unchecked(*p.get_unchecked(s[i] as usize) as usize) // / 2.; // } // } else { // for i in 0..VOCAB { // for j in 0..VOCAB { // lp += cnts[i][j] as f32 // * TRANS // .get_unchecked(p[j] as usize) // .get_unchecked(p[i] as usize); // } // } // } // } lp } fn weights(s: &[u8], p: &[u8; VOCAB], ys: &Tensor) -> Vec<f64> { // dbg!(to_string(&permute(s, p))); let mut swaps = [[0.; VOCAB]; VOCAB]; let mut cnts = [[0; VOCAB]; VOCAB]; for i in 0..s.len() { for j in 0..VOCAB { // NOTE: ys is padded let v = ys.i((i as i64, j as i64)).double_value(&[]); swaps[p[s[i] as usize] as usize][j] += v; cnts[p[s[i] as usize] as usize][j] += 1; } } let mut w = vec![0.; VOCAB * VOCAB]; // print!(" "); // for i in 0..VOCAB { // print!("{:5} ", p // .iter() // .position(|x| *x as usize == i) // .unwrap()); // } // println!(); for i in 0..VOCAB { // print!("{:2} ", p // .iter() // .position(|x| *x as usize == i) // .unwrap()); for j in 0..VOCAB { if j < i { if cnts[i][j] > 0 { w[i * VOCAB + j] += swaps[i][j] / cnts[i][j] as f64; } if cnts[j][i] > 0 { w[i * VOCAB + j] += swaps[j][i] / cnts[j][i] as f64; } // print!("{:>5} ", swaps[i][j] / cnts[i][j] as f64); // Oh no floating point comparison spooky spooky if w[i * VOCAB + j] != 0. { // TODO: adjust this factor w[i * VOCAB + j] = (w[i * VOCAB + j] / 2.5).exp(); } // print!("{:5.2} ", w[i * VOCAB + j]); } } // println!(); } w } fn refiner(net: &impl Module, s: &[u8], p2: &[u8; VOCAB], iters: usize) -> ([u8; VOCAB], f64) { let mut p = *p2; // Cleanup using CNN let mut rng = Xoshiro256PlusPlus::from_rng(&mut rand::rng()); // Use padding to properly handle the ends let (mut ys, mut loss) = probs_padded(net, &permute(s, &p)); // dbg!(loss, to_string(&permute(s, &p))); if loss > 3. { // Not worth trying return (p, loss); } let mut pbest = *p2; let mut lossbest = loss; let mut w = weights(s, &p, &ys); for _ in 0..300 * iters / s.len() { let dist = WeightedIndex::new(&w).unwrap(); let sample = dist.sample(&mut rng); let a = p .iter() .position(|x| *x as usize == sample / VOCAB) .unwrap(); let b = p .iter() .position(|x| *x as usize == sample % VOCAB) .unwrap(); let mut q = p; (q[a], q[b]) = (q[b], q[a]); let (ys2, loss2) = probs_padded(net, &permute(s, &q)); let acc = f64::min(0., (loss - loss2) * 20.); if rng.random::<f64>() < acc.exp() { ys = ys2; p = q; loss = loss2; w = weights(s, &p, &ys); } if loss2 < lossbest { lossbest = loss2; pbest = q; } } // dbg!(lossbest, to_string(&permute(s, &pbest))); (pbest, lossbest) } fn decode( net: &impl Module, s: &[u8], grams: &[[[f32; VOCAB]; VOCAB]; VOCAB], ) -> ([u8; VOCAB], f64) { // Initialize using naive freqs let mut porig = [0; VOCAB]; let mut cnt = [0; VOCAB]; for c in s { cnt[*c as usize] += 1; } let mut indices: Vec<_> = (0..VOCAB).collect(); indices.sort_unstable_by_key(|&a| cnt[a]); for i in 0..VOCAB { porig[indices[i]] = VOCAB_ORDER[i]; } // MCMC for some iters let mut cnts = [[0; VOCAB]; VOCAB]; for i in 1..s.len() { cnts[s[i - 1] as usize][s[i] as usize] += 1; } // Is this easy to decode? let mut rng = Xoshiro256PlusPlus::from_rng(&mut rand::rng()); for _ in 0..20 { let mut p = porig; let mut lp = logprob(s, &p, &cnts, grams); for _ in 0..10000 { let a = rng.random_range(0..VOCAB); let b = rng.random_range(0..VOCAB); if a != b { let mut q = p; (q[a], q[b]) = (q[b], q[a]); let lp2 = logprob(s, &q, &cnts, grams); let acc = f32::min(0., lp2 - lp); if rng.random::<f32>() < acc.exp() { p = q; lp = lp2; } } } let (p2, loss) = refiner(net, s, &p, 100); // dbg!(to_string(&permute(s, &p2)), loss); if loss < 0.7 { // Probably done return refiner(net, s, &p2, 1000); } } // Nope we're gonna have to try hard ugh // FEARLESS CONCURRENCY let real_iters = ITERS * 100 / s.len() / THREADS * THREADS; let mut res: Vec<([u8; VOCAB], f32)> = vec![]; thread::scope(|sc| { let mut handles = vec![]; for _ in 0..THREADS { let handle = sc.spawn(move || { let mut tmp = vec![]; let mut rng = Xoshiro256PlusPlus::from_rng(&mut rand::rng()); for _k in 0..real_iters / THREADS { let mut p = porig; let mut lp = logprob(s, &p, &cnts, grams); let mut pbest = p; let mut lpbest = lp; // TODO: investigate if we need more/less iters here for _ in 0..10000 { let a = rng.random_range(0..VOCAB); let b = rng.random_range(0..VOCAB); if a != b { let mut q = p; (q[a], q[b]) = (q[b], q[a]); // We *could* reuse some of the computation from lp // But logprob is already a tight nested loop so it's not much faster // And floating point garbage builds up let lp2 = logprob(s, &q, &cnts, grams); let acc = f32::min(0., lp2 - lp); if rng.random::<f32>() < acc.exp() { p = q; lp = lp2; if lp2 > lpbest { lpbest = lp2; pbest = q; } } } } // dbg!(_k, to_string(&permute(s, &pbest)), lpbest); tmp.push((pbest, lpbest)); } tmp }); handles.push(handle); } for handle in handles { res.extend(&handle.join().unwrap()); } }); let mut ord = (0..real_iters).collect::<Vec<_>>(); ord.sort_unstable_by(|&i, &j| res[j].1.partial_cmp(&res[i].1).unwrap()); // Try 50 best perms let mut pbest = porig; let mut lossbest = 100.; // Don't retry perms that decode to the same thing let mut tried = HashSet::new(); let mut i = 0; while tried.len() < 40 { let p2 = res[ord[i]].0; let s2 = permute(s, &p2); if !tried.contains(&s2) { tried.insert(s2); let (p, loss) = refiner(net, s, &p2, 50); if loss < lossbest { pbest = p; lossbest = loss; } if lossbest < 0.7 { // Probably done break; } } i += 1; } refiner(net, s, &pbest, 1000) } fn finish( net: &impl Module, text: &[u8], pl: &[u8; 28], pr: &[u8; 28], jbest: usize, ) -> (Vec<u8>, f64) { let pl2 = &refiner(net, &text[..jbest], pl, 1000).0; let pr2 = &refiner(net, &text[jbest..], pr, 1000).0; let ans = [permute(&text[..jbest], pl2), permute(&text[jbest..], pr2)].concat(); (ans.clone(), probs(net, &ans).1) } pub fn main() { let device = Device::cuda_if_available(); // dbg!(device); let mut vs = nn::VarStore::new(device); let net = net(&vs.root()); let args: Vec<String> = env::args().collect(); // dbg!(args.len()); if args.len() == 3 { no_grad(|| { // Load trigrams let file = File::open("trigrams").unwrap(); let mut reader = BufReader::new(file); let mut grams = [[[0.; VOCAB]; VOCAB]; VOCAB]; for i in 0..VOCAB { for j in 0..VOCAB { for k in 0..VOCAB { let mut buf = String::new(); reader.read_line(&mut buf).unwrap(); buf.pop(); // Remove newline grams[i][j][k] = buf.parse().unwrap(); } } } // Decode vs.load("model.safetensors").unwrap(); let text: Vec<_> = args[1].chars().map(char_to_label).collect(); if args[2].to_lowercase() != "true" { // No breakpoint // dbg!(probs(&net, &text).1); // return; let p = decode(&net, &text, &grams).0; println!("{}", to_string(&permute(&text, &p))); } else { // Decode left and right halves let m = text.len() / 2; let (pl, lossl) = decode(&net, &text[..m], &grams); let (pr, lossr) = decode(&net, &text[m..], &grams); // dbg!(to_string(&permute(&text[..m], &pl))); // dbg!(to_string(&permute(&text[m..], &pr))); // dbg!(lossl, lossr); let mut reallossbest = 100.; let mut ans = vec![]; // let mut should_loop = true; // This is spaghetti code yeah I know if lossl < 1.5 { // First half correct let mut deque = VecDeque::new(); let mut sum = 0.; let ys = probs(&net, &permute(&text[m - WINDOW..], &pl)).0; for i in m - WINDOW + KERNEL as usize..text.len() - KERNEL as usize { deque.push_back( -ys.i(( (i + WINDOW - m) as i64 - KERNEL, pl[text[i] as usize] as i64, )) .double_value(&[]), ); sum += deque.back().unwrap(); if deque.len() > WINDOW { sum -= deque.front().unwrap(); deque.pop_front(); } // dbg!(i, sum / WINDOW as f64); if deque.len() == WINDOW && sum / WINDOW as f64 > 4. { // Breakpoint probably in i - WINDOW to i - 10 let pr2 = decode(&net, &text[i - 10..], &grams).0; let mut jbest = 0; let mut lossbest = 100.; let mut pt = permute( &text[i - WINDOW - KERNEL as usize..i + KERNEL as usize], &pr2, ); for j in i - WINDOW..i { pt[j + WINDOW - i + KERNEL as usize] = pl[text[j] as usize]; let loss = probs(&net, &pt).1; // dbg!(j, to_string(&pt), loss); if loss < lossbest { jbest = j + 1; // Off-by-1 error lossbest = loss; } } // dbg!(jbest, i); (ans, reallossbest) = finish(&net, &text, &pl, &pr2, jbest); // should_loop = false; break; } } if reallossbest == 100. { // No breakpoint found??? ans = permute(&text, &refiner(&net, &text, &pl, 1000).0); reallossbest = probs(&net, &ans).1; } } if lossl >= 1.5 || lossr < 1.5 { // Second half correct // Same alg, just reversed let mut deque = VecDeque::new(); let mut sum = 0.; let ys = probs(&net, &permute(&text[..m + WINDOW], &pr)).0; // TODO: Are these indices correct??? for i in (KERNEL as usize..m + WINDOW - KERNEL as usize).rev() { deque.push_back( -ys.i((i as i64 - KERNEL, pr[text[i] as usize] as i64)) .double_value(&[]), ); sum += deque.back().unwrap(); if deque.len() > WINDOW { sum -= deque.front().unwrap(); deque.pop_front(); } if deque.len() == WINDOW && sum / WINDOW as f64 > 4. { // Breakpoint probably in i + 10 to i + WINDOW let pl2 = decode(&net, &text[..i + 10], &grams).0; let mut jbest = 0; let mut lossbest = 100.; let mut pt = permute( &text[i - KERNEL as usize..i + WINDOW + KERNEL as usize], &pr, ); for j in i..i + WINDOW { pt[j + KERNEL as usize - i] = pl2[text[j] as usize]; let loss = probs(&net, &pt).1; // dbg!(j, to_string(&pt), loss); if loss < lossbest { jbest = j + 1; // Off-by-1 error lossbest = loss; } } // dbg!(jbest, i); let (ans2, loss2) = finish(&net, &text, &pl2, &pr, jbest); if loss2 < reallossbest { ans = ans2; reallossbest = loss2; } // should_loop = false; break; } } if reallossbest == 100. { // No breakpoint found??? ans = permute(&text, &refiner(&net, &text, &pr, 1000).0); // reallossbest = probs(&net, &ans).1; } } // if should_loop { // // failed to find breakpoint // sleep(time::Duration::from_secs(1000000)); // } println!("{}", to_string(&ans)); } }) } else { // vs.load("model.safetensors").unwrap(); // Train model let mut opt = nn::Adam::default().build(&vs, LR).unwrap(); opt.set_weight_decay(WD); let file = File::open("wikitext").unwrap(); let reader = BufReader::new(file); let mut data = vec![]; for line in reader.lines() { // net doesn't accept u8 as input, must be i64 let linedata: Vec<_> = line .unwrap() .chars() .map(|c| char_to_label(c) as i64) .collect(); for i in 0..linedata.len() / SEQ_LEN { // Pad with spaces // I should probably use a dedicated padding token but I don't want to retrain data.push( [ &[26; KERNEL as usize], &linedata[i * SEQ_LEN..(i + 1) * SEQ_LEN], &[26; KERNEL as usize], ] .concat(), ); } } for i in 0..EPOCHS { for j in 0..data.len() / BS { let xs = Tensor::from_slice2(&data[j * BS..(j + 1) * BS]).to(device); let ys = net.forward(&xs).transpose(-2, -1); let loss = ys.nll_loss_nd::<Tensor>( &xs.i((.., KERNEL..SEQ_LEN as i64 + KERNEL)), None, Reduction::Mean, -100, ); println!("{i} {j} {loss}"); opt.backward_step(&loss); } } vs.save("model.safetensors").unwrap(); // vs.freeze(); // let mut closure = |input: &[Tensor]| vec![net.forward(&input[0])]; // let model = CModule::create_by_tracing( // "MyModule", // "forward", // &[Tensor::zeros([1, 784], (tch::Kind::Int64, device))], // &mut closure, // ) // .unwrap(); // // I think this has the input size hardcoded though sad // model.save("model.pt").unwrap(); } }
-