CogView2 / patch
hysts's picture
hysts HF staff
Copy files from https://github.com/hysts/CogView2_demo
9c00f5c
raw
history blame
1.75 kB
diff --git a/coglm_strategy.py b/coglm_strategy.py
index cba87ce..40e4ece 100755
--- a/coglm_strategy.py
+++ b/coglm_strategy.py
@@ -8,6 +8,7 @@
# here put the import lib
import os
+import pathlib
import sys
import math
import random
@@ -57,7 +58,8 @@ class CoglmStrategy:
self._is_done = False
self.outlier_count_down = 5
self.vis_list = [[]for i in range(16)]
- self.cluster_labels = torch.tensor(np.load('cluster_label.npy'), device='cuda', dtype=torch.long)
+ cluster_label_path = pathlib.Path(__file__).parent / 'cluster_label.npy'
+ self.cluster_labels = torch.tensor(np.load(cluster_label_path), device='cuda', dtype=torch.long)
self.top_k_cluster = top_k_cluster
@property
@@ -91,4 +93,4 @@ class CoglmStrategy:
def finalize(self, tokens, mems):
self._is_done = False
- return tokens, mems
\ No newline at end of file
+ return tokens, mems
diff --git a/sr_pipeline/dsr_sampling.py b/sr_pipeline/dsr_sampling.py
index a0d0298..f721573 100755
--- a/sr_pipeline/dsr_sampling.py
+++ b/sr_pipeline/dsr_sampling.py
@@ -8,6 +8,7 @@
# here put the import lib
import os
+import pathlib
import sys
import math
import random
@@ -27,7 +28,8 @@ class IterativeEntfilterStrategy:
self.invalid_slices = invalid_slices
self.temperature = temperature
self.topk = topk
- self.cluster_labels = torch.tensor(np.load('cluster_label.npy'), device='cuda', dtype=torch.long)
+ cluster_label_path = pathlib.Path(__file__).parents[1] / 'cluster_label.npy'
+ self.cluster_labels = torch.tensor(np.load(cluster_label_path), device='cuda', dtype=torch.long)
self.temperature2 = temperature2