Skip to content
Projects
Groups
Snippets
Help
Loading...
Help
Support
Keyboard shortcuts
?
Submit feedback
Sign in / Register
Toggle navigation
H
Hydra Node Http
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
Hydra Node Http
Commits
7ae5e142
Commit
7ae5e142
authored
Oct 01, 2022
by
kurumuz
Browse files
Options
Browse Files
Download
Email Patches
Plain Diff
multi-knn
parent
fb3c04d0
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
18 additions
and
14 deletions
+18
-14
hydra_node/models.py
hydra_node/models.py
+18
-14
No files found.
hydra_node/models.py
View file @
7ae5e142
...
...
@@ -709,14 +709,18 @@ class EmbedderModel(nn.Module):
import
requests
knn_folder
=
config
.
knn_folder
self
.
model
=
SentenceTransformer
(
'sentence-transformers/all-MiniLM-L6-v2'
)
.
cuda
()
self
.
index
=
pickle
.
load
(
requests
.
get
(
f
"{knn_folder}/index.pkl"
,
stream
=
'True'
)
.
raw
)
self
.
tag_count
=
pickle
.
load
(
requests
.
get
(
f
"{knn_folder}/all_tags.pkl"
,
stream
=
'True'
)
.
raw
)
r
=
requests
.
get
(
f
"{knn_folder}/knn.index"
,
stream
=
'True'
)
with
open
(
"knn.index"
,
"wb"
)
as
f
:
f
.
write
(
r
.
content
)
self
.
knn
=
faiss
.
read_index
(
"knn.index"
)
self
.
tag_count_sorted
=
sorted
(
self
.
tag_count
.
items
(),
key
=
lambda
x
:
x
[
1
],
reverse
=
True
)
self
.
indexes
=
{}
for
folder
in
knn_folder
.
split
(
","
):
name
,
url
=
folder
.
split
(
":"
)
index
=
pickle
.
load
(
requests
.
get
(
f
"{url}/index.pkl"
,
stream
=
'True'
)
.
raw
)
tag_count
=
pickle
.
load
(
requests
.
get
(
f
"{url}/all_tags.pkl"
,
stream
=
'True'
)
.
raw
)
tag_count_sorted
=
sorted
(
tag_count
.
items
(),
key
=
lambda
x
:
x
[
1
],
reverse
=
True
)
r
=
requests
.
get
(
f
"{url}/knn.index"
,
stream
=
'True'
)
with
open
(
"knn.index"
,
"wb"
)
as
f
:
f
.
write
(
r
.
content
)
knn
=
faiss
.
read_index
(
"knn.index"
)
self
.
indexes
[
name
]
=
[
index
,
tag_count
,
tag_count_sorted
,
knn
]
def
__call__
(
self
,
sentences
):
with
torch
.
no_grad
():
...
...
@@ -725,9 +729,11 @@ class EmbedderModel(nn.Module):
def
get_top_k
(
self
,
request
):
text
=
request
.
prompt
model
=
request
.
model
index
,
tag_count
,
tag_count_sorted
,
knn
=
self
.
indexes
[
model
]
#check if text is a substring in tag_count.keys()
found
=
[]
for
tag
,
count
in
self
.
tag_count_sorted
:
for
tag
,
count
in
tag_count_sorted
:
if
len
(
tag
)
>=
len
(
text
)
and
tag
.
startswith
(
text
):
found
.
append
([
tag
,
count
,
0
])
...
...
@@ -735,11 +741,11 @@ class EmbedderModel(nn.Module):
embedding
=
self
([
text
])
#print(embedding.dtype)
k
=
15
D
,
I
=
self
.
knn
.
search
(
embedding
,
k
)
D
,
I
=
knn
.
search
(
embedding
,
k
)
D
,
I
=
D
.
squeeze
(),
I
.
squeeze
()
for
i
,
id
in
enumerate
(
I
):
tag
=
self
.
index
[
id
]
count
=
self
.
tag_count
[
tag
]
tag
=
index
[
id
]
count
=
tag_count
[
tag
]
prob
=
D
[
i
]
results
.
append
([
tag
,
count
,
prob
])
...
...
@@ -753,8 +759,6 @@ class EmbedderModel(nn.Module):
#filter results for >0.5 confidence unless it has the search text in it and confidence is >0.4
results
=
[
x
for
x
in
results
if
x
[
2
]
>
0.5
or
(
x
[
2
]
>
0.4
and
text
in
x
[
0
])]
found
=
sorted
(
found
,
key
=
lambda
x
:
x
[
1
],
reverse
=
True
)
print
(
found
)
print
(
results
)
if
len
(
found
)
>
0
:
#results = results[:-len(found)]
results
=
found
+
results
...
...
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