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
Files changed (1) hide show
  1. app.py +8 -4
app.py CHANGED
@@ -41,10 +41,7 @@ def add_to_index(files, index):
41
 
42
 
43
  @spaces.GPU
44
- def query_index(query, index):
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