Skip to content
Projects
Groups
Snippets
Help
Loading...
Help
Support
Keyboard shortcuts
?
Submit feedback
Sign in / Register
Toggle navigation
B
Basedformer
Project overview
Project overview
Details
Activity
Releases
Repository
Repository
Files
Commits
Branches
Tags
Contributors
Graph
Compare
Locked Files
Issues
0
Issues
0
List
Boards
Labels
Service Desk
Milestones
Merge Requests
0
Merge Requests
0
CI / CD
CI / CD
Pipelines
Jobs
Schedules
Security & Compliance
Security & Compliance
Dependency List
License Compliance
Packages
Packages
List
Container Registry
Analytics
Analytics
CI / CD
Code Review
Insights
Issues
Repository
Value Stream
Wiki
Wiki
Snippets
Snippets
Members
Members
Collapse sidebar
Close sidebar
Activity
Graph
Create a new issue
Jobs
Commits
Issue Boards
Open sidebar
novelai-storage
Basedformer
Commits
a1556cf7
Commit
a1556cf7
authored
Apr 26, 2022
by
novelailab
Browse files
Options
Browse Files
Download
Email Patches
Plain Diff
push
parent
1f8d04b2
Changes
2
Hide whitespace changes
Inline
Side-by-side
Showing
2 changed files
with
90 additions
and
15 deletions
+90
-15
basedformer/sampling.py
basedformer/sampling.py
+3
-0
scripts/test_cache.py
scripts/test_cache.py
+87
-15
No files found.
basedformer/sampling.py
0 → 100644
View file @
a1556cf7
import
torch
import
sys
scripts/test_cache.py
View file @
a1556cf7
...
@@ -13,6 +13,82 @@ def print_top_k(logits, tokenizer, k):
...
@@ -13,6 +13,82 @@ def print_top_k(logits, tokenizer, k):
for
token
in
topk_ind
[
x
,
y
,
:]
.
tolist
():
for
token
in
topk_ind
[
x
,
y
,
:]
.
tolist
():
print
(
tokenizer
.
decode
([
token
]),
end
=
" | "
)
print
(
tokenizer
.
decode
([
token
]),
end
=
" | "
)
def
apply_top_k
(
logits
,
k
):
# filter the logits that are not in the top-k to -inf
# keep top_k_ind and filter the rest
top_k_values
=
logits
.
topk
(
k
)[
0
]
remove_mask
=
logits
<
top_k_values
[:,
-
1
]
.
unsqueeze
(
0
)
logits
[
remove_mask
==
True
]
=
-
float
(
"inf"
)
return
logits
def
apply_top_p
(
logits
,
p
):
sorted
,
indices
=
torch
.
sort
(
logits
,
descending
=
True
)
cumulative_probs
=
torch
.
cumsum
(
sorted
,
dim
=-
1
)
cumulative_probs
=
cumulative_probs
.
scatter
(
dim
=-
1
,
index
=
indices
,
src
=
cumulative_probs
)
remove_mask
=
cumulative_probs
>
p
logits
[
remove_mask
==
True
]
=
-
float
(
"inf"
)
return
logits
def
apply_tfs
(
logits
,
tfs
):
sorted
,
indices
=
torch
.
sort
(
logits
,
descending
=
True
)
d
=
sorted
d
=
d
[:,
1
:]
-
d
[:,
:
-
1
]
d
=
d
[:,
1
:]
-
d
[:,
:
-
1
]
d
=
d
.
abs
()
d
=
d
/
d
.
sum
(
dim
=-
1
)
.
view
(
1
,
-
1
)
.
T
cumulative_probs
=
torch
.
cumsum
(
d
,
dim
=-
1
)
cumulative_probs
=
cumulative_probs
.
scatter
(
dim
=-
1
,
index
=
indices
,
src
=
cumulative_probs
)
remove_mask
=
cumulative_probs
>
tfs
logits
[
remove_mask
==
True
]
=
-
float
(
"inf"
)
return
logits
def
temperature
(
logits
,
temperature
):
logits
=
logits
/
temperature
return
logits
def
generate
(
forward
,
prompt_tokens
,
temperature
,
tokens_to_generate
=
50
,
ops_list
=
[{
"temp"
:
0.9
}]):
with
torch
.
no_grad
():
in_tokens
=
prompt_tokens
kv
=
None
tokens_generated
=
[]
soft_required
=
[
"top_k"
,
"top_p"
]
op_map
=
{
"top_k"
:
apply_top_k
,
"top_p"
:
apply_top_p
,
"temp"
:
temperature
,
"tfs"
:
apply_tfs
}
for
_
in
range
(
tokens_to_generate
):
logits
,
kv
=
forward
(
in_tokens
,
cache
=
True
,
kv
=
kv
)
# always work on softmax logits to make sure all models
# behave similarly as logprobs can be quite different
# TODO: can break compatibility with novelai presets.
# logits should be the last token in the sequence
logits
=
logits
[:,
-
1
,
:]
logits
=
torch
.
log_softmax
(
logits
,
dim
=-
1
)
#can save one softmax here by not applying softmax for the first op,
#need to take the softmax out of the necessary functions though
for
i
,
ops
in
enumerate
(
ops_list
):
batch
=
[]
for
op
,
value
in
ops
.
items
():
if
op
in
soft_required
:
item
=
torch
.
log_softmax
(
logits
[
i
,
:,
:],
dim
=-
1
)
item
=
op_map
[
op
](
item
,
value
)
batch
.
append
(
item
)
logits
=
torch
.
cat
(
batch
,
dim
=
0
)
logits
=
torch
.
softmax
(
logits
,
dim
=-
1
)
logits
=
torch
.
multinomial
(
logits
,
1
)
in_tokens
=
logits
tokens_generated
.
append
(
logits
)
tokens_generated
=
torch
.
cat
(
tokens_generated
,
dim
=-
1
)
return
tokens_generated
def
main
():
def
main
():
tokenizer
=
AutoTokenizer
.
from_pretrained
(
'gpt2'
)
tokenizer
=
AutoTokenizer
.
from_pretrained
(
'gpt2'
)
prompt
=
"""I fucked her with my huge donut, when she seen my donut she went"""
prompt
=
"""I fucked her with my huge donut, when she seen my donut she went"""
...
@@ -28,20 +104,16 @@ def main():
...
@@ -28,20 +104,16 @@ def main():
ic
(
time
.
perf_counter
()
-
t
)
ic
(
time
.
perf_counter
()
-
t
)
with
torch
.
no_grad
():
ops
=
{
kv
=
None
"top_k"
:
40
,
tokens_to_generate
=
50
"top_p"
:
0.9
,
in_tokens
=
tokens
"temp"
:
0.9
,
accum_tokens
=
[]
}
for
x
in
range
(
tokens_to_generate
):
logits
,
kv
=
model
(
in_tokens
,
cache
=
True
,
kv
=
kv
)
tokens_generated
=
generate
(
model
.
forward
,
tokens
,
40
,
ops
=
ops
)
in_tokens
=
logits
[:,
-
1
,
:]
.
topk
(
1
)[
1
]
tokens_generated
=
tokenizer
.
decode
(
tokens_generated
.
squeeze
()
.
tolist
())
#in_tokens = torch.cat([in_tokens, logits[:, -1, :].topk(1)[1]], dim=1)
ic
(
prompt
)
print
(
tokenizer
.
decode
(
in_tokens
.
squeeze
(
1
)
.
tolist
()[
-
1
]),
end
=
" | "
)
ic
(
tokens_generated
)
#accum_tokens = torch.cat(accum_tokens, dim=1)
#accum_tokens = accum_tokens.squeeze(0).tolist()
#print("\n Final token list")
#print(tokenizer.decode(accum_tokens))
if
__name__
==
"__main__"
:
if
__name__
==
"__main__"
:
main
()
main
()
\ No newline at end of file
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment