Spaces:
Running
on
Zero
Running
on
Zero
Antoine Chaffin
commited on
Commit
·
3f46c94
1
Parent(s):
346feeb
Using the GPU only when generating the answer to avoid pickling the index
Browse files
app.py
CHANGED
@@ -41,10 +41,7 @@ def add_to_index(files, index):
|
|
41 |
|
42 |
|
43 |
@spaces.GPU
|
44 |
-
def
|
45 |
-
res = index(query, k=1)
|
46 |
-
retrieved_image = res["documents"][0][0]["image"]
|
47 |
-
|
48 |
messages = [
|
49 |
{
|
50 |
"role": "user",
|
@@ -80,6 +77,13 @@ def query_index(query, index):
|
|
80 |
skip_special_tokens=True,
|
81 |
clean_up_tokenization_spaces=False,
|
82 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
83 |
|
84 |
return output_text[0], retrieved_image
|
85 |
|
|
|
41 |
|
42 |
|
43 |
@spaces.GPU
|
44 |
+
def generate_answer(query, retrieved_image):
|
|
|
|
|
|
|
45 |
messages = [
|
46 |
{
|
47 |
"role": "user",
|
|
|
77 |
skip_special_tokens=True,
|
78 |
clean_up_tokenization_spaces=False,
|
79 |
)
|
80 |
+
return output_text[0]
|
81 |
+
|
82 |
+
|
83 |
+
def query_index(query, index):
|
84 |
+
res = index(query, k=1)
|
85 |
+
retrieved_image = res["documents"][0][0]["image"]
|
86 |
+
output_text = generate_answer(query, retrieved_image)
|
87 |
|
88 |
return output_text[0], retrieved_image
|
89 |
|