Spaces:
Build error
Build error
File size: 1,677 Bytes
452b173 |
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 |
#include "rep_penalty.h"
#include <cstdlib>
#include <cstring>
void rep_penalty_cpu
(
const int vocab_size,
const uint64_t* sequence,
float* rep_mask,
const float penalty_max,
const int sustain,
const int decay,
const int seq_len
)
{
float v = penalty_max;
float dv = decay ? (1.0f - penalty_max) / (float) decay : 0.0f;
int s = sustain == -1 ? seq_len : sustain;
int beg = seq_len - s - decay;
if (beg < 0) beg = 0;
for (int i = 0; i < vocab_size; i++) rep_mask[i] = 1.0f;
for (int i = seq_len; i > beg;)
{
uint64_t t = sequence[--i];
if (v > rep_mask[t]) rep_mask[t] = v;
if (--s < 0) v += dv;
}
}
bool* g_rep_mask = NULL;
int g_vocab_size = 0;
void apply_rep_penalty_cpu
(
const int vocab_size,
const uint64_t* sequence,
const float penalty_max,
const int sustain,
const int decay,
const int seq_len,
float* logits
)
{
if (vocab_size != g_vocab_size)
{
if (g_rep_mask) free(g_rep_mask);
g_vocab_size = vocab_size;
g_rep_mask = (bool*) malloc(g_vocab_size * sizeof(bool));
}
memset(g_rep_mask, 0, g_vocab_size * sizeof(bool));
float v = penalty_max;
float dv = decay ? (1.0f - penalty_max) / (float) decay : 0.0f;
int s = sustain == -1 ? seq_len : sustain;
int beg = seq_len - s - decay;
if (beg < 0) beg = 0;
for (int i = seq_len; i > beg;)
{
uint64_t t = sequence[--i];
if (!g_rep_mask[t])
{
if (logits[t] > 0.0) logits[t] /= v;
else logits[t] *= v;
g_rep_mask[t] = true;
}
if (--s < 0) v += dv;
}
} |