Skip to content

Commit

Permalink
feat: 判题
Browse files Browse the repository at this point in the history
Signed-off-by: YdrMaster <[email protected]>
  • Loading branch information
YdrMaster committed Jan 17, 2025
1 parent 99d583d commit 103f42c
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 9 deletions.
1 change: 1 addition & 0 deletions models/llama/nvidia-gpu/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -244,6 +244,7 @@ impl<'blk> Weights<'blk> {
load! {
attn_norm
attn_qkv
attn_qkv_bias
attn_o
ffn_norm
ffn_gate_inp
Expand Down
41 changes: 32 additions & 9 deletions test-utils/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -145,27 +145,50 @@ pub fn test_infer(
let maybe_file = Path::new(prompt);
if maybe_file.is_file() {
let file = std::fs::read_to_string(maybe_file).unwrap();
let mut correct = 0;
let mut total = 0;
for line in file.lines() {
let line = serde_json::from_str::<String>(line).unwrap();
let prompt = format!("<s>{line}\n");

// print_now!("{prompt}");
let line =
serde_json::from_str::<serde_json::Map<String, serde_json::Value>>(line).unwrap();
let serde_json::Value::String(prompt) = &line["origin_prompt"] else {
unreachable!()
};
let serde_json::Value::String(gold) = &line["gold"] else {
unreachable!()
};

let mut tokens = tokenizer.encode(&prompt);
let mut ans = String::new();
let mut pos = 0;
let mut result = false;
for _ in 0..max_steps {
let next = lm(&tokens, pos);

pos += tokens.len();
if next == eos {
break;
}

let piece = tokenizer.decode(next);
print_now!("{piece}");
pos += tokens.len();
tokens = vec![next];

let piece = tokenizer.decode(next);
ans.push_str(&piece);

if let Some(x) = piece.chars().find(|c| matches!(c, 'A'..='D')) {
result = x == gold.chars().next().unwrap();
break;
}
}
println!()

if result {
correct += 1
}
total += 1;

println!(
"({correct:>4}/{total:<4} {:6.2}%) {} {gold} {ans}",
100. * (correct as f64 / total as f64),
if result { "✔" } else { "✘" },
)
}
return;
}
Expand Down

0 comments on commit 103f42c

Please sign in to comment.