Core ML
legal
citation-verification
hallucination-detection
apple-neural-engine
hand-written-weights
citenet-coreml / bench /CiteScan.swift
alexwengg's picture
CiteNet v1.5: legal-citation checker with hand-written weights (Core ML)
42b4a6a verified
Raw History Blame Contribute Delete
13.6 kB
// Native citation scanner + encoder: a byte-level port of citenet/text.py
// (CITE_RE, normalize, encode). scanparity.swift checks it against Python.
import Foundation
let SEQ = 64
let MAXCHARS = 46
@inline(__always) func isDigit(_ b: UInt8) -> Bool { b >= 0x30 && b <= 0x39 }
@inline(__always) func isUpper(_ b: UInt8) -> Bool { b >= 0x41 && b <= 0x5A }
@inline(__always) func isLower(_ b: UInt8) -> Bool { b >= 0x61 && b <= 0x7A }
@inline(__always) func isAlpha(_ b: UInt8) -> Bool { isUpper(b) || isLower(b) }
struct Text {
let p: UnsafePointer<UInt8>
let n: Int
@inline(__always) subscript(_ i: Int) -> UInt8 { i < n ? p[i] : 0 }
/// Length in bytes of the whitespace character at i (Python str `\s`), or 0.
@inline(__always) func ws(_ i: Int) -> Int {
let b = self[i]
if b == 0x20 || (b >= 0x09 && b <= 0x0D) || (b >= 0x1C && b <= 0x1F) { return i < n ? 1 : 0 }
if b < 0x80 { return 0 }
let b1 = self[i + 1], b2 = self[i + 2]
if b == 0xC2 && (b1 == 0x85 || b1 == 0xA0) { return 2 }
if b == 0xE1 && b1 == 0x9A && b2 == 0x80 { return 3 }
if b == 0xE2 && b1 == 0x80 && ((b2 >= 0x80 && b2 <= 0x8A) || b2 == 0xA8 || b2 == 0xA9 || b2 == 0xAF) { return 3 }
if b == 0xE2 && b1 == 0x81 && b2 == 0x9F { return 3 }
if b == 0xE3 && b1 == 0x80 && b2 == 0x80 { return 3 }
return 0
}
@inline(__always) func skipWS(_ i: Int) -> Int { var j = i; while true { let w = ws(j); if w == 0 { return j }; j += w } }
@inline(__always) func isRSQ(_ i: Int) -> Bool { self[i] == 0xE2 && self[i + 1] == 0x80 && self[i + 2] == 0x99 }
@inline(__always) func isEnDash(_ i: Int) -> Bool { self[i] == 0xE2 && self[i + 1] == 0x80 && self[i + 2] == 0x93 }
/// Python `\w` for the character ending just before i (ASCII exact; non-ASCII letters count as word chars).
func wordBefore(_ i: Int) -> Bool {
if i <= 0 { return false }
let b = p[i - 1]
if b < 0x80 { return isAlpha(b) || isDigit(b) || b == 0x5F }
var s = i - 1
while s > 0 && (p[s] & 0xC0) == 0x80 { s -= 1 }
guard let sc = String(bytes: UnsafeBufferPointer(start: p + s, count: i - s), encoding: .utf8)?.unicodeScalars.first else { return false }
return sc.properties.isAlphabetic || sc.properties.numericType != nil
}
/// Python `\w` for the character starting at i.
func wordAt(_ i: Int) -> Bool {
if i >= n { return false }
let b = p[i]
if b < 0x80 { return isAlpha(b) || isDigit(b) || b == 0x5F }
var e = i + 1
while e < n && (p[e] & 0xC0) == 0x80 { e += 1 }
return wordBefore(e)
}
/// TOKEN = [A-Z][A-Za-z'’&]*\.? | \d+(?:d|th|st|nd|rd)\b -> end index or -1
@inline(__always) func token(_ i: Int) -> Int {
let b = self[i]
if isUpper(b) {
var q = i + 1
while q < n {
let c = p[q]
if isAlpha(c) || c == 0x27 || c == 0x26 { q += 1 } else if isRSQ(q) { q += 3 } else { break }
}
if self[q] == 0x2E { q += 1 }
return q
}
if isDigit(b) {
var q = i
while isDigit(self[q]) { q += 1 }
let c0 = self[q], c1 = self[q + 1]
var e = -1
if c0 == 0x64 { e = q + 1 } // d
else if (c0 == 0x74 && c1 == 0x68) || (c0 == 0x73 && c1 == 0x74) ||
(c0 == 0x6E && c1 == 0x64) || (c0 == 0x72 && c1 == 0x64) { e = q + 2 } // th st nd rd
return (e >= 0 && !wordAt(e)) ? e : -1
}
return -1
}
/// (?:at\s+)?\*?\d{1,8} -> end or -1
@inline(__always) func starDigits(_ s: Int, _ m: Int) -> Int {
var t = s
if self[t] == 0x2A { t += 1 }
let d0 = t
while isDigit(self[t]) && t - d0 < m { t += 1 }
return t > d0 ? t : -1
}
@inline(__always) func pageAt(_ i: Int, maxDigits: Int) -> Int {
if self[i] == 0x61 && self[i + 1] == 0x74 {
let k = skipWS(i + 2)
if k > i + 2 { let r = starDigits(k, maxDigits); if r >= 0 { return r } }
}
return starDigits(i, maxDigits)
}
/// One CITE_RE match starting at i (i must be at a digit). Returns (end, accept) or nil;
/// accept = the reporter has an uppercase letter and does not open with a STOP word.
func match(_ i: Int) -> (Int, Bool)? {
var j = i
while isDigit(self[j]) { j += 1 }
if j - i > 4 { return nil }
let k = skipWS(j)
if k == j { return nil }
// greedy maximal tokens, at most 5, separated by at most one whitespace char
var ends = [Int](); ends.reserveCapacity(5)
var upper = [Bool](); upper.reserveCapacity(5)
var t = k
while ends.count < 5 {
let e = token(t)
if e < 0 { break }
ends.append(e); upper.append(isUpper(self[t]) || (upper.last ?? false))
t = e + ws(e) // \s? : at most one whitespace char
}
if ends.isEmpty { return nil }
for c in stride(from: ends.count - 1, through: 0, by: -1) {
let e = ends[c]
let s = skipWS(e)
var pg = -1
if self[s] == 0x2C { // \s*,\s*at\s+ page ("976 F. 2d, at 540")
let q = skipWS(s + 1)
if self[q] == 0x61 && self[q + 1] == 0x74 {
let r = skipWS(q + 2)
if r > q + 2 { pg = starDigits(r, 8) }
}
} else if s > e { // \s+(?:at\s+)? page
pg = pageAt(s, maxDigits: 8)
}
if pg < 0 { continue }
var end = pg
// pincites: (?:\s*,\s*(?:at\s+)?\*?\d+(?:\s*[-–]\s*\*?\d+)?)*
while true {
var q = skipWS(end)
if self[q] != 0x2C { break }
q = skipWS(q + 1)
let d = pageAt(q, maxDigits: Int.max)
if d < 0 { break }
var e2 = d
var r = skipWS(d)
if self[r] == 0x2D || isEnDash(r) {
r = skipWS(r + (self[r] == 0x2D ? 1 : 3))
if self[r] == 0x2A { r += 1 }
let d0 = r
while isDigit(self[r]) { r += 1 }
if r > d0 { e2 = r }
}
end = e2
}
// paren: \s*\([^()]{0,60}?\d{4}\)
let po = skipWS(end)
if self[po] == 0x28 {
var q = po + 1, chars = 0
while q < n && p[q] != 0x28 && p[q] != 0x29 && chars <= 64 {
if (p[q] & 0xC0) != 0x80 { chars += 1 }
q += 1
}
if q < n && p[q] == 0x29 && chars >= 4 && chars - 4 <= 60 &&
isDigit(p[q - 1]) && isDigit(p[q - 2]) && isDigit(p[q - 3]) && isDigit(p[q - 4]) {
end = q + 1
}
}
let first = Array(UnsafeBufferPointer(start: p + k, count: ends[0] - k))
return (end, upper[c] && !STOP.contains(first))
}
return nil
}
/// finditer over [lo, hi): calls emit(start, end) for each accepted citation.
func scan(_ lo: Int, _ hi: Int, _ emit: (Int, Int) -> Void) {
var i = lo
while i < hi {
if isDigit(p[i]) && !wordBefore(i) {
if let (e, up) = match(i) {
if up { emit(i, e) }
i = e
continue
}
while i < hi && isDigit(p[i]) { i += 1 } // \b fails inside the rest of this run
continue
}
i += 1
}
}
}
/// citenet/text.py STOP: words that open a match but never a reporter
let STOP: Set<[UInt8]> = Set("""
No. Nos. No Id. Id Ibid. Section Sec. In On At Page Pages Rule Rules Cir. Supp. Tr. Year Years
East West North South And Or The Of To CA Art. Article Para. Count Counts Exhibit Ex. Doc. Dkt. ECF
January February March April May June July August September October November December
Jan. Feb. Mar. Apr. Jun. Jul. Aug. Sep. Sept. Oct. Nov. Dec.
C Civ Civ. Cr Cr. Crim. Crim CV CR
""".split(whereSeparator: { $0 == " " || $0 == "\n" }).map { Array($0.utf8) })
// ---------------------------------------------------------------- normalize + encode
let OTHERB: UInt8 = 0xFF
/// Map the match bytes to one byte per character (normalize() steps 1-3: NFKC-ish,
/// quote/dash/section mapping, whitespace collapse + strip, lowercase).
@inline(__always) func charsOf(_ t: Text, _ s: Int, _ e: Int, into out: inout [UInt8]) {
out.removeAll(keepingCapacity: true)
var i = s
var pendingSpace = false
while i < e {
let w = t.ws(i)
if w > 0 { pendingSpace = !out.isEmpty; i += w; continue }
var c: UInt8
let b = t.p[i]
if b < 0x80 {
c = isUpper(b) ? b + 32 : (b == 0x60 ? 0x27 : b); i += 1
} else {
let b1 = t[i + 1], b2 = t[i + 2]
if b == 0xE2 && b1 == 0x80 && (b2 == 0x98 || b2 == 0x99) { c = 0x27; i += 3 }
else if b == 0xE2 && b1 == 0x80 && (b2 == 0x93 || b2 == 0x94) { c = 0x2D; i += 3 }
else if b == 0xC2 && b1 == 0xA7 { c = 0x73; i += 2 }
else {
c = OTHERB
i += 1
while i < e && (t.p[i] & 0xC0) == 0x80 { i += 1 }
}
}
if pendingSpace { out.append(0x20); pendingSpace = false }
out.append(c)
}
}
@inline(__always) func isWordB(_ c: UInt8) -> Bool { isLower(c) || isDigit(c) || c == 0x5F || c == OTHERB }
let EXPL: [[UInt8]] = ["quoting", "citing", "internal", "alteration", "emphasis", "quotation", "affirmed", "aff'd",
"rev'd", "cert", "vacated", "overruled", "abrogated", "superseded", "modified"].map { Array($0.utf8) }
/// _EXPL: \s*\((?:quoting|citing|...|\d+ ) -> cut the string at the leftmost match
@inline(__always) func cutExplanatory(_ s: inout [UInt8]) {
let n = s.count
var i = 0
while i < n {
if s[i] == 0x28 {
var hit = false
for w in EXPL where i + 1 + w.count <= n && s[(i + 1)..<(i + 1 + w.count)].elementsEqual(w) { hit = true; break }
if !hit {
var q = i + 1
while q < n && isDigit(s[q]) { q += 1 }
if q > i + 1 && q < n && s[q] == 0x20 { hit = true }
}
if hit {
var cut = i
while cut > 0 && s[cut - 1] == 0x20 { cut -= 1 }
s.removeSubrange(cut..<n)
return
}
}
i += 1
}
}
func normalizeChars(_ s: inout [UInt8], _ tmp: inout [UInt8]) {
cutExplanatory(&s)
// _TRAIL: (?<=[\d)])[\s.,;:]+$ | [\s,;:]+$ (leftmost position where either matches)
var k = s.count
while k > 0 && (s[k - 1] == 0x20 || s[k - 1] == 0x2E || s[k - 1] == 0x2C || s[k - 1] == 0x3B || s[k - 1] == 0x3A) { k -= 1 }
if k < s.count {
var cut = -1
for pos in k..<s.count {
let prev = pos > 0 ? s[pos - 1] : 0
if isDigit(prev) || prev == 0x29 { cut = pos; break }
var allNoDot = true
for q in pos..<s.count where s[q] == 0x2E { allNoDot = false; break }
if allNoDot { cut = pos; break }
}
if cut >= 0 { s.removeSubrange(cut..<s.count) }
}
// _JOIN: \.\s+(?=(?!at\b)[a-z]|\d+(?:d|th|st|nd|rd)\b) -> "."
tmp.removeAll(keepingCapacity: true)
var i = 0
let n = s.count
while i < n {
let c = s[i]
tmp.append(c)
if c == 0x2E && i + 1 < n && s[i + 1] == 0x20 {
let j = i + 2
var join = false
if j < n && isLower(s[j]) {
let isAt = s[j] == 0x61 && j + 1 < n && s[j + 1] == 0x74 && !(j + 2 < n && isWordB(s[j + 2]))
join = !isAt
} else if j < n && isDigit(s[j]) {
var q = j
while q < n && isDigit(s[q]) { q += 1 }
var e = -1
if q < n && s[q] == 0x64 { e = q + 1 }
else if q + 1 < n && ((s[q] == 0x74 && s[q + 1] == 0x68) || (s[q] == 0x73 && s[q + 1] == 0x74) ||
(s[q] == 0x6E && s[q + 1] == 0x64) || (s[q] == 0x72 && s[q + 1] == 0x64)) { e = q + 2 }
if e >= 0 && !(e < n && isWordB(s[e])) { join = true }
}
if join { i += 2; continue }
}
i += 1
}
swap(&s, &tmp)
if s.count > MAXCHARS {
let head = Array(s[0..<24]), tail = Array(s[(s.count - (MAXCHARS - 24))...])
s = head + tail
}
}
let CODE: [UInt8] = {
var t = [UInt8](repeating: 47, count: 256)
for d in 0..<10 { t[0x30 + d] = UInt8(3 + d) }
for l in 0..<26 { t[0x61 + l] = UInt8(13 + l) }
t[0x20] = 39; t[0x2E] = 40; t[0x2C] = 41; t[0x28] = 42; t[0x29] = 43; t[0x27] = 44; t[0x26] = 45; t[0x2D] = 46
return t
}()
/// Write BOS + codes + EOS + PAD (SEQ Float16 values) for citation bytes [s, e).
@inline(__always) func encode(_ t: Text, _ s: Int, _ e: Int, _ dst: UnsafeMutablePointer<Float16>,
_ buf: inout [UInt8], _ tmp: inout [UInt8]) {
charsOf(t, s, e, into: &buf)
normalizeChars(&buf, &tmp)
dst[0] = 1
var k = 1
for c in buf where k < SEQ - 1 { dst[k] = Float16(CODE[Int(c)]); k += 1 }
dst[k] = 2
k += 1
while k < SEQ { dst[k] = 0; k += 1 }
}