StevenChen16 commited on
Commit
52c56bf
β€’
1 Parent(s): d61cc55

third commit - add spaces.GPU

Browse files
Files changed (1) hide show
  1. app.py +3 -0
app.py CHANGED
@@ -36,6 +36,8 @@ from langchain_huggingface import HuggingFaceEmbeddings
36
  from tqdm import tqdm
37
  import joblib
38
 
 
 
39
  # Token streamer for generation
40
  class TokenStreamer(BaseStreamer):
41
  def __init__(self, skip_prompt: bool = False, timeout=None):
@@ -170,6 +172,7 @@ def query_vector_store(vector_store: FAISS, query, k=4, relevance_threshold=0.8)
170
  context = [doc.page_content for doc in similar_docs]
171
  return context
172
 
 
173
  class ModelWorker:
174
  def __init__(self, model_path, device='cuda'):
175
  self.device = device
 
36
  from tqdm import tqdm
37
  import joblib
38
 
39
+ import spaces
40
+
41
  # Token streamer for generation
42
  class TokenStreamer(BaseStreamer):
43
  def __init__(self, skip_prompt: bool = False, timeout=None):
 
172
  context = [doc.page_content for doc in similar_docs]
173
  return context
174
 
175
+ @spaces.GPU
176
  class ModelWorker:
177
  def __init__(self, model_path, device='cuda'):
178
  self.device = device