209 lines
6.9 KiB
C#
209 lines
6.9 KiB
C#
using System;
|
|
using System.Collections.Generic;
|
|
using LLama.Extensions;
|
|
using LLama.Native;
|
|
|
|
namespace LLama.Sampling;
|
|
|
|
/// <summary>
|
|
/// An implementation of ISamplePipeline which mimics the default llama.cpp sampling
|
|
/// </summary>
|
|
public sealed class DefaultSamplingPipeline
|
|
: BaseSamplingPipeline
|
|
{
|
|
/// <summary>
|
|
/// Bias values to add to certain logits
|
|
/// </summary>
|
|
public Dictionary<int, float> LogitBias { get; } = new();
|
|
|
|
/// <summary>
|
|
/// Repetition penalty, as described in https://arxiv.org/abs/1909.05858
|
|
/// </summary>
|
|
public float RepeatPenalty { get; set; }
|
|
|
|
/// <summary>
|
|
/// Frequency penalty as described by OpenAI: https://platform.openai.com/docs/api-reference/chat/create<br />
|
|
/// Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing frequency in the text
|
|
/// so far, decreasing the model's likelihood to repeat the same line verbatim.
|
|
/// </summary>
|
|
public float AlphaFrequency
|
|
{
|
|
get => _alphaFreq;
|
|
set
|
|
{
|
|
if (value < -2)
|
|
throw new ArgumentOutOfRangeException(nameof(value), "AlphaFrequency must be greater than -2");
|
|
if (value > 2)
|
|
throw new ArgumentOutOfRangeException(nameof(value), "AlphaFrequency must be less than 2");
|
|
_alphaFreq = value;
|
|
}
|
|
}
|
|
private float _alphaFreq;
|
|
|
|
/// <summary>
|
|
/// Presence penalty as described by OpenAI: https://platform.openai.com/docs/api-reference/chat/create<br />
|
|
/// Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in the
|
|
/// text so far, increasing the model's likelihood to talk about new topics.
|
|
/// </summary>
|
|
public float AlphaPresence
|
|
{
|
|
get => _alphaPresence;
|
|
set
|
|
{
|
|
if (value < -2)
|
|
throw new ArgumentOutOfRangeException(nameof(value), "AlphaFrequency must be greater than -2");
|
|
if (value > 2)
|
|
throw new ArgumentOutOfRangeException(nameof(value), "AlphaFrequency must be less than 2");
|
|
_alphaPresence = value;
|
|
}
|
|
}
|
|
private float _alphaPresence;
|
|
|
|
/// <summary>
|
|
/// Temperature to apply (higher temperature is more "creative")
|
|
/// </summary>
|
|
public float Temperature { get; set; } = 0.75f;
|
|
|
|
/// <summary>
|
|
/// Number of tokens to keep in TopK sampling
|
|
/// </summary>
|
|
public int TopK { get; set; }
|
|
|
|
/// <summary>
|
|
/// Z value for tail free sampling
|
|
/// </summary>
|
|
public float TailFreeZ { get; set; }
|
|
|
|
/// <summary>
|
|
/// P value for locally typical sampling
|
|
/// </summary>
|
|
public float TypicalP { get; set; }
|
|
|
|
/// <summary>
|
|
/// P value for TopP sampling
|
|
/// </summary>
|
|
public float TopP { get; set; } = 1f;
|
|
|
|
/// <summary>
|
|
/// P value for MinP sampling
|
|
/// </summary>
|
|
public float MinP { get; set; }
|
|
|
|
/// <summary>
|
|
/// Whether the newline value should be protected from being modified by logit bias and repeat penalty
|
|
/// </summary>
|
|
public bool PenalizeNewline { get; set; } = false;
|
|
|
|
/// <inheritdoc />
|
|
protected override void ProcessLogits(SafeLLamaContextHandle ctx, Span<float> logits, ReadOnlySpan<LLamaToken> lastTokens)
|
|
{
|
|
// Apply logit bias
|
|
foreach (var (key, value) in LogitBias)
|
|
logits[key] += value;
|
|
|
|
}
|
|
|
|
/// <inheritdoc />
|
|
protected override LLamaToken ProcessTokenDataArray(SafeLLamaContextHandle ctx, LLamaTokenDataArray candidates, ReadOnlySpan<LLamaToken> lastTokens)
|
|
{
|
|
// Only apply repetition penalty if we really must. Otherwise avoid all this work
|
|
if (lastTokens.Length > 0 && (RepeatPenalty != 0 || AlphaFrequency != 0 || AlphaPresence != 0))
|
|
{
|
|
// Save the logit value for the newline token
|
|
var (nlIndex, nlLogit) = PenalizeNewline ? GetNewlineLogit(ctx, candidates) : (-1, 0);
|
|
|
|
// Apply penalties to candidates
|
|
candidates.RepetitionPenalty(ctx, lastTokens, RepeatPenalty, AlphaFrequency, AlphaPresence);
|
|
|
|
// Restore newline token
|
|
if (!PenalizeNewline)
|
|
SetNewlineLogit(ctx, candidates, nlIndex, nlLogit);
|
|
}
|
|
|
|
// Apply the normal llama.cpp pipeline
|
|
candidates.ApplyGrammar(ctx, Grammar);
|
|
candidates.TopK(ctx, TopK);
|
|
candidates.TailFree(ctx, TailFreeZ);
|
|
candidates.LocallyTypical(ctx, TypicalP);
|
|
candidates.TopP(ctx, TopP);
|
|
candidates.MinP(ctx, MinP);
|
|
candidates.Temperature(ctx, Temperature);
|
|
return candidates.SampleToken(ctx);
|
|
}
|
|
|
|
private static (int, float) GetNewlineLogit(SafeLLamaContextHandle ctx, LLamaTokenDataArray candidates)
|
|
{
|
|
var nlToken = ctx.ModelHandle.Tokens.Newline;
|
|
|
|
if (nlToken.HasValue)
|
|
{
|
|
// Try using the ID as an index
|
|
if (candidates.data.Span[(int)nlToken].id == nlToken)
|
|
return ((int)nlToken, candidates.data.Span[(int)nlToken].logit);
|
|
|
|
// Exhaustive search
|
|
var span = candidates.data.Span;
|
|
for (var i = 0; i < span.Length; i++)
|
|
{
|
|
if (span[i].id == nlToken)
|
|
return (i, span[i].logit);
|
|
}
|
|
}
|
|
|
|
return (-1, 0);
|
|
}
|
|
|
|
private static void SetNewlineLogit(SafeLLamaContextHandle ctx, LLamaTokenDataArray candidates, int indexHint, float logit)
|
|
{
|
|
var nlToken = ctx.ModelHandle.Tokens.Newline;
|
|
if (!nlToken.HasValue)
|
|
return;
|
|
|
|
// Try checking the index where we found it last time. It might not be there if `RepetitionPenalty` changed order
|
|
if (indexHint >= 0 && candidates.data.Span[indexHint].id == nlToken)
|
|
{
|
|
candidates.data.Span[indexHint].logit = logit;
|
|
return;
|
|
}
|
|
|
|
// Didn't find it, do an exhaustive search for it
|
|
var span = candidates.data.Span;
|
|
for (var i = 0; i < candidates.data.Length; i++)
|
|
{
|
|
if (span[i].id == nlToken)
|
|
{
|
|
span[i].logit = logit;
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
|
|
/// <inheritdoc />
|
|
public override void Accept(SafeLLamaContextHandle ctx, LLamaToken token)
|
|
{
|
|
Grammar?.AcceptToken(ctx, token);
|
|
}
|
|
|
|
/// <inheritdoc />
|
|
public override ISamplingPipeline Clone()
|
|
{
|
|
var clone = new DefaultSamplingPipeline();
|
|
|
|
foreach (var (k, v) in LogitBias)
|
|
clone.LogitBias.Add(k, v);
|
|
|
|
clone.Grammar = Grammar?.Clone();
|
|
clone.RepeatPenalty = RepeatPenalty;
|
|
clone.AlphaFrequency = AlphaFrequency;
|
|
clone.AlphaPresence = AlphaPresence;
|
|
clone.Temperature = Temperature;
|
|
clone.TopK = TopK;
|
|
clone.TailFreeZ = TailFreeZ;
|
|
clone.TypicalP = TypicalP;
|
|
clone.TopP = TopP;
|
|
clone.MinP = MinP;
|
|
clone.PenalizeNewline = PenalizeNewline;
|
|
|
|
return clone;
|
|
}
|
|
} |