File size: 6,883 Bytes
fb62ba4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
// elephant-1 ONNX demo: chat generation + one end-to-end scored decision.
//
//   dotnet run -- --model ../model-cuda-int4 --cuda --prompt "What is the capital of France?"
//   dotnet run -- --model ../model-cuda-int4 --cuda --score-demo
//
// Tokenizer: gigatoken_cs on the build's tokenizer.json — verified byte-identical
// to HuggingFace transformers on 19 texts (unicode, CJK, emoji, specials).
// Chat bytes below mirror the checkpoint's chat template for a single user turn
// (verified against apply_chat_template output).
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;

static class Program
{
    const string SystemPrompt =
        "Apply the question to the state. Choose exactly one of the listed options. " +
        "Respond with only its uppercase letter, with no explanation or reasoning.";

    static int Main(string[] args)
    {
        string modelDir = Arg(args, "--model", "../model-cuda-int4");
        bool cuda = Has(args, "--cuda");
        string prompt = Arg(args, "--prompt", "What is the capital of France? Answer with the city name only.");
        bool scoreDemo = Has(args, "--score-demo");
        int maxNewTokens = int.Parse(Arg(args, "--max-new-tokens", "512"));
        // The build dirs carry the GenAI-fixed tokenizer (ByteLevel-only), which
        // gigatoken_cs 0.12.0 cannot load; the pristine original tokenizes
        // byte-identically (verified 19/19), so load that instead.
        string tokPath = Arg(args, "--tokenizer", "../orig_tokenizer/tokenizer.json");

        var tok = Gigatoken.Tokenizer.FromJson(File.ReadAllText(tokPath));
        using (tok)
        using (var lm = new OnnxLm(Path.Combine(modelDir, "model.onnx"), cuda))
        {
            if (scoreDemo) return ScoreDemo(tok, lm);
            Generate(tok, lm, modelDir, prompt, maxNewTokens);
            return 0;
        }
    }

    // ---- chat generation (causal decode on the bidi graph) ----
    static void Generate(Gigatoken.Tokenizer tok, OnnxLm lm, string modelDir, string prompt, int maxNewTokens)
    {
        string chat = "<|im_start|>user\n" + prompt + "<|im_end|>\n" +
                      "<|im_start|>assistant\n<think>\n\n</think>\n\n";
        long[] ids = tok.Encode(chat).Select(i => (long)i).ToArray();
        Console.WriteLine($"  prompt tokens: {ids.Length}");
        var eos = EosIds(modelDir);
        var t0 = DateTime.UtcNow;
        var outIds = lm.Generate(ids, maxNewTokens, eos, t =>
        {
            Console.Write(tok.DecodeToString(new[] { (uint)t }));
            Console.Out.Flush();
        });
        Console.WriteLine($"\n  tokens={outIds.Count} total={(DateTime.UtcNow - t0).TotalSeconds:F1}s");
    }

    static HashSet<long> EosIds(string modelDir)
    {
        var set = new HashSet<long>();
        try
        {
            var doc = JsonDocument.Parse(File.ReadAllText(Path.Combine(modelDir, "genai_config.json")));
            var eos = doc.RootElement.GetProperty("model").GetProperty("eos_token_id");
            if (eos.ValueKind == JsonValueKind.Array)
                foreach (var e in eos.EnumerateArray()) set.Add(e.GetInt64());
            else set.Add(eos.GetInt64());
        }
        catch { set.Add(248046); set.Add(248044); }
        return set;
    }

    // ---- one scored decision: support.json's refund question ----
    // Prompt rendering mirrors score_elephant.py (qev-labels-v2) for a
    // criterion-less noul question; the program prints the prompt SHA256 so
    // you can check byte-equality against the Python renderer.
    static int ScoreDemo(Gigatoken.Tokenizer tok, OnnxLm lm)
    {
        string state = "Customer: my order arrived broken, I want my money back.";
        string instructions = "Should the agent offer a refund?";
        string t = "the statement is true / the answer is yes";
        string f = "the statement is false / the answer is no";
        string suffix = $"Question: {instructions}\n\nOptions:\nA. yes: {t}\nB. no: {f}\n\nAnswer with one letter: A, B.";
        string user = $"<state>\n{state}\n</state>\n\n" + suffix;
        string full = "<|im_start|>system\n" + SystemPrompt + "<|im_end|>\n" +
                      "<|im_start|>user\n" + user + "<|im_end|>\n" +
                      "<|im_start|>assistant\n<think>\n\n</think>\n\n";
        Console.WriteLine("prompt sha256: " + Sha(full));

        long[] ids = tok.Encode(full).Select(i => (long)i).ToArray();
        int prefix = CommonPrefix(tok, full, $"<state>\n{state}\n</state>\n\n");
        Console.WriteLine($"branch tokens: {ids.Length}, prefix tokens: {prefix}");

        var bias = new float[ids.Length * ids.Length];
        for (int i = 0; i < ids.Length; i++)
            for (int j = 0; j < ids.Length; j++)
                bias[i * ids.Length + j] = (j < prefix || i >= prefix) ? 0f : OnnxLm.NegInf;

        uint[] slots = { tok.Encode("A")[0], tok.Encode("B")[0] };
        if (tok.Encode("A").Length != 1 || tok.Encode("B").Length != 1)
        { Console.Error.WriteLine("slot labels are not single tokens"); return 1; }

        var (last, st) = lm.Prefill(ids, bias);
        using (st)
        {
            double[] z = { last[slots[0]], last[slots[1]] };
            double m = Math.Max(z[0], z[1]);
            double[] p = { Math.Exp(z[0] - m), Math.Exp(z[1] - m) };
            double s = p[0] + p[1];
            Console.WriteLine($"refund noul: {p[0] / s:F4} (Python: 0.8791 CUDA int8, 0.9342 CUDA int4, 0.9069 CPU, 0.8921 CPU int8)");
            return 0;
        }
    }

    // Shared-prefix length: tokenize every branch fully, then find the longest
    // common token prefix with the state-block text (marker method, mirrors
    // qev _prefix_and_suffix_ids: the shared prefix ends at the state block).
    static int CommonPrefix(Gigatoken.Tokenizer tok, string full, string stateBlock)
    {
        long[] fullIds = tok.Encode(full).Select(i => (long)i).ToArray();
        const string mark = "QEV_PREFIX_MARK";
        string head = "<|im_start|>system\n" + SystemPrompt + "<|im_end|>\n" +
                      "<|im_start|>user\n" + stateBlock + mark;
        long[] headIds = tok.Encode(head).Select(i => (long)i).ToArray();
        int n = Math.Min(headIds.Length, fullIds.Length);
        while (n > 0 && !fullIds.Take(n).SequenceEqual(headIds.Take(n))) n--;
        return n;
    }

    static string Sha(string s)
    {
        var h = SHA256.HashData(Encoding.UTF8.GetBytes(s));
        return BitConverter.ToString(h).Replace("-", "").ToLowerInvariant();
    }

    static string Arg(string[] a, string name, string dflt)
    {
        for (int i = 0; i + 1 < a.Length; i++) if (a[i] == name) return a[i + 1];
        return dflt;
    }

    static bool Has(string[] a, string name) => a.Contains(name);
}