| import Foundation | |
| /// Length prior over candidate segment lengths. Port of `create_prior_function` | |
| /// in `wtpsplit/utils/priors.py`. A prior returns a non-negative weight for a | |
| /// given segment length; `0` forbids that length. | |
| public enum Prior: String, Sendable, CaseIterable { | |
| case uniform | |
| case gaussian | |
| case clippedPolynomial | |
| } | |
| enum PriorFactory { | |
| /// Build the base prior closure. `targetLength`/`spread` are ignored by | |
| /// `.uniform` (matching the reference, which omits them for uniform). | |
| static func make(_ kind: Prior, maxLength: Int, targetLength: Int, spread: Int) -> (Int) -> Double { | |
| let maxL = maxLength | |
| switch kind { | |
| case .uniform: | |
| return { length in length > maxL ? 0.0 : 1.0 } | |
| case .gaussian: | |
| let target = Double(targetLength) | |
| let sp = Double(spread) | |
| return { length in | |
| if length > maxL { return 0.0 } | |
| let z = (Double(length) - target) / sp | |
| return exp(-0.5 * z * z) | |
| } | |
| case .clippedPolynomial: | |
| let target = Double(targetLength) | |
| let falloff = 1.0 / (Double(spread) * Double(spread)) | |
| return { length in | |
| if length > maxL { return 0.0 } | |
| let d = Double(length) - target | |
| return max(1.0 - falloff * d * d, 0.0) | |
| } | |
| } | |
| } | |
| /// Wrap a base prior with the soft-overflow decay used by | |
| /// `run_segmentation.py`: lengths past `softCap` are damped by a Gaussian | |
| /// tail so plain word gaps can't exploit the slack, while a strong pause can | |
| /// still pull a break into the overflow zone. `overflow == 0` ⇒ hard cap. | |
| static func withOverflow(_ base: @escaping (Int) -> Double, softCap: Int, overflow: Int) -> (Int) -> Double { | |
| guard overflow > 0 else { return base } | |
| let decay = Double(overflow) | |
| return { length in | |
| let b = base(length) | |
| if length <= softCap { return b } | |
| let t = (Double(length - softCap)) / decay | |
| return b * exp(-(t * t)) | |
| } | |
| } | |
| } | |