antirez / antirez/LLM-FTC-sampling

I wrote a `transformers` LogitProcessor implementation

Aperta
#2 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub
Lingua principale
Python
Stelle
30
Fork
1
Metriche di merge delle PR
Nessuna PR unita negli ultimi 30g

Descrizione

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.

Guida per i contributori

Nessuna guida per i contributori indicizzata per questo repository

Valutazione

Questa issue non è ancora stata valutata.

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.