Spaces:
Runtime error
Runtime error
StevenChen16
commited on
Commit
β’
52c56bf
1
Parent(s):
d61cc55
third commit - add spaces.GPU
Browse files
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
|