antirez / antirez/LLM-FTC-sampling
I wrote a `transformers` LogitProcessor implementation
- 主要語言
- Python
- 星號
- 30
- 分支
- 1
- PR 合併指標
- 30 天內沒有已合併 PR
描述
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.
貢獻指南
這個儲存庫沒有索引到貢獻指南
評估
這個 Issue 還沒有評估資料。