File size: 5,574 Bytes
26d5b81 | 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 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | #include "neuroflow/generative_model.hpp"
#include <algorithm>
#include <cmath>
#include <iostream>
namespace neuroflow {
GenerativeModel::GenerativeModel(const CausalLMConfig& lm_config,
std::unique_ptr<Tokenizer> tokenizer,
NeuroFlowModel* nf_model)
: tokenizer_(std::move(tokenizer)), neuroflow_model_(nf_model) {
lm_head_ = std::make_unique<CausalLMHead>(lm_config);
sampler_ = std::make_unique<TopKSampling>();
}
GenerateOutput GenerativeModel::generate(const std::string& prompt, const GenerateConfig& config) {
GenerateOutput output;
std::vector<size_t> input_ids = tokenizer_->encode(prompt);
std::vector<size_t> generated_ids;
size_t seed = config.random_seed;
if (seed == 0) {
std::random_device rd;
seed = rd();
}
std::mt19937 rng(seed);
lm_head_->clear_cache();
if (input_ids.empty()) {
output.finish_reason = FinishReason::GEN_ERROR;
return output;
}
Tensor logits;
if (input_ids.size() > 1) {
std::vector<size_t> prefix(input_ids.begin(), input_ids.end() - 1);
logits = lm_head_->forward(prefix);
}
size_t last_id = input_ids.back();
for (size_t step = 0; step < config.max_new_tokens; ++step) {
size_t pos = input_ids.size() - 1 + step;
logits = lm_head_->forward_step(last_id, pos);
logits = apply_sn_gating(lm_head_->last_hidden_, logits);
logits = inject_memory(lm_head_->last_hidden_, logits);
logits = apply_repetition_penalty(logits, config, generated_ids);
logits = apply_punct_penalty(logits, config);
Tensor probs = sampler_->apply(logits, config, generated_ids);
size_t next_id = sampler_->sample(probs, rng);
generated_ids.push_back(next_id);
if (next_id == config.eos_id) {
output.finish_reason = FinishReason::EOS_TOKEN;
break;
}
last_id = next_id;
}
if (output.finish_reason != FinishReason::EOS_TOKEN) {
output.finish_reason = FinishReason::MAX_LENGTH;
}
output.token_ids = generated_ids;
output.text = tokenizer_->decode(generated_ids);
output.cache_stats = lm_head_->cache_stats();
return output;
}
Tensor GenerativeModel::apply_sn_gating(const Tensor& hidden, const Tensor& logits) {
if (!neuroflow_model_) return logits;
try {
auto sn_out = neuroflow_model_->sn->forward(hidden);
float ecn_gate = sn_out.gates.as_fp32()[0];
float dmn_gate = sn_out.gates.as_fp32()[1];
if (std::isnan(ecn_gate) || std::isnan(dmn_gate) ||
(ecn_gate == 0.0f && dmn_gate == 0.0f)) {
std::cerr << "Warning: SN gate NaN/zero, falling back to equal weighting" << std::endl;
ecn_gate = 0.5f;
dmn_gate = 0.5f;
}
return logits;
} catch (...) {
return logits;
}
}
Tensor GenerativeModel::inject_memory(const Tensor& query, const Tensor& logits) {
if (!neuroflow_model_) return logits;
try {
auto mem_out = neuroflow_model_->memory->retrieve(query);
const float* mp = mem_out.retrieved.as_fp32();
bool all_zero = true;
for (size_t i = 0; i < mem_out.retrieved.numel(); ++i) {
if (std::abs(mp[i]) > 1e-8f) { all_zero = false; break; }
}
if (all_zero) return logits;
return logits;
} catch (...) {
return logits;
}
}
Tensor GenerativeModel::apply_repetition_penalty(Tensor logits, const GenerateConfig& config,
const std::vector<size_t>& generated) {
if (config.repetition_penalty <= 1.0f || generated.empty()) return logits;
float* data = logits.as_fp32();
size_t n = logits.numel();
size_t window = std::min(generated.size(), static_cast<size_t>(20));
for (size_t i = generated.size() - window; i < generated.size(); ++i) {
size_t tid = generated[i];
if (tid < n) {
if (data[tid] > 0) {
data[tid] /= config.repetition_penalty;
} else {
data[tid] *= config.repetition_penalty;
}
}
}
return logits;
}
Tensor GenerativeModel::apply_punct_penalty(Tensor logits, const GenerateConfig& config) {
if (config.punct_penalty <= 0.0f || config.punct_ids.empty()) return logits;
float* data = logits.as_fp32();
size_t n = logits.numel();
for (size_t pid : config.punct_ids) {
if (pid < n) {
data[pid] -= config.punct_penalty;
}
}
return logits;
}
void GenerativeModel::set_strategy(SamplingStrategyType type) {
switch (type) {
case SamplingStrategyType::GREEDY:
sampler_ = std::make_unique<GreedyDecoding>();
break;
case SamplingStrategyType::TOP_K:
sampler_ = std::make_unique<TopKSampling>();
break;
case SamplingStrategyType::TOP_P:
sampler_ = std::make_unique<TopPSampling>();
break;
case SamplingStrategyType::TOP_K_TOP_P:
sampler_ = std::make_unique<TopKTopPSampling>();
break;
}
}
void GenerativeModel::clear_cache() { lm_head_->clear_cache(); }
CacheStats GenerativeModel::cache_stats() const { return lm_head_->cache_stats(); }
} // namespace neuroflow |