Update app.py
Browse files
app.py
CHANGED
@@ -14,9 +14,8 @@ chunk_size = int(os.environ.get("CHUNK_SIZE", 128))
|
|
14 |
default_max_characters = int(os.environ.get("DEFAULT_MAX_CHARACTERS", 258))
|
15 |
|
16 |
model = SentenceTransformer(model_name)
|
17 |
-
|
18 |
|
19 |
-
@spaces.GPU
|
20 |
def embed(queries, chunks) -> dict[str, list[tuple[str, float]]]:
|
21 |
query_embeddings = model.encode(queries, prompt_name="query")
|
22 |
document_embeddings = model.encode(chunks)
|
|
|
14 |
default_max_characters = int(os.environ.get("DEFAULT_MAX_CHARACTERS", 258))
|
15 |
|
16 |
model = SentenceTransformer(model_name)
|
17 |
+
model.to(device="cuda")
|
18 |
|
|
|
19 |
def embed(queries, chunks) -> dict[str, list[tuple[str, float]]]:
|
20 |
query_embeddings = model.encode(queries, prompt_name="query")
|
21 |
document_embeddings = model.encode(chunks)
|