tangger/Search-R1
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

clean up retrieval cache

Browse Source
This commit is contained in:
PeterGriffinJin
2025-03-23 14:33:14 +00:00
parent 6272082a64
commit f5204213d3
1 changed files with 8 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
search_r1/search/retrieval_server.py
View File

@@ -125,6 +125,10 @@ class Encoder:
query_emb = query_emb.detach().cpu().numpy() query_emb = query_emb.detach().cpu().numpy()
query_emb = query_emb.astype(np.float32, order="C") query_emb = query_emb.astype(np.float32, order="C")
del inputs, output
torch.cuda.empty_cache()
return query_emb return query_emb
class BaseRetriever: class BaseRetriever:
@@ -266,6 +270,10 @@ class DenseRetriever(BaseRetriever):
results.extend(batch_results) results.extend(batch_results)
scores.extend(batch_scores) scores.extend(batch_scores)
del batch_emb, batch_scores, batch_idxs, query_batch, flat_idxs, batch_results
torch.cuda.empty_cache()
if return_score: if return_score:
return results, scores return results, scores
else: else: