antirez / antirez/LLM-FTC-sampling

I wrote a `transformers` LogitProcessor implementation

Aberta
#2 0 comentários 0 reações 0 responsáveis Ver no GitHub
Linguagem predominante
Python
Estrelas
30
Forks
1
Métricas de merge de PRs
Nenhum PR com merge em 30d

Descrição

bellissima idea!

It would be nice to have a logit processor with this sampling algorithm in HuggingFace's `transformers`. I wrote one pretty quickly:
```py
from transformers import LogitsProcessor
import numpy as np
import torch

class FirstTokenCutoffLogitsProcessor(LogitsProcessor):
def __init__(self, cutoff=0.7):
self.cutoff = cutoff

def __call__(self, _, scores):
for i in range(len(scores)):
scores_i = scores[i]
logits_i = torch.softmax(scores_i, dim=-1)
np_scores_i = np.array(logits_i)
sorted_indices = np.argsort(np_scores_i)
sorted_indices = sorted_indices[::-1]
j = 1
t0 = np_scores_i[sorted_indices[0]]
while j < len(np_scores_i) and 1 - (np_scores_i[sorted_indices[j]] / t0) < self.cutoff:
j += 1
accepted_logits = []
for k in range(0, j):
accepted_logits.append(float(np_scores_i[sorted_indices[k]]))
accepted_logits = np.array(accepted_logits)
idx = torch.multinomial(torch.tensor(
accepted_logits), num_samples=1)
token_id = sorted_indices[idx]
scores[i, token_id] = -1e9
return scores
```

Here is how one can use it:
```py
from transformers import AutoTokenizer, AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M")
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M")

prompt = "Hello, I'm a language model,"
input_ids = tokenizer.encode(prompt, return_tensors="pt")

processor = FirstTokenCutoffLogitsProcessor()
outs = model.generate(
input_ids, max_new_tokens=120, do_sample=False, logits_processor=[processor]
)[0]
print(tokenizer.decode(outs))
```

I can make a PR to put this in if you'd like.

Guia de contribuição

Nenhum guia de contribuição indexado para este repositório

Avaliação

Esta issue ainda não foi avaliada.

Receba novas issues na sua caixa de entrada

Um resumo curto de issues do GitHub para quem está começando.