Skip to content

Commit

Permalink
Add new method: trace
Browse files Browse the repository at this point in the history
  • Loading branch information
ignorejjj committed Jul 6, 2024
1 parent 9fcfa0d commit b9d6fa1
Show file tree
Hide file tree
Showing 3 changed files with 6 additions and 5 deletions.
1 change: 1 addition & 0 deletions examples/methods/run_exp.py
Original file line number Diff line number Diff line change
Expand Up @@ -429,6 +429,7 @@ def trace(args):
config_dict = {'save_note': save_note,
'gpu_id':args.gpu_id,
'dataset_name':args.dataset_name,
'refiner_name': 'kg-trace',
'trace_config': trace_config}

# preparation
Expand Down
6 changes: 3 additions & 3 deletions flashrag/refiner/kg_refiner.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from flashrag.retriever.encoder import Encoder, STEncoder
from flashrag.utils import hash_object

class KGRefiner(BaseRefiner):
class KGTraceRefiner(BaseRefiner):
def __init__(self, config, retriever=None, generator=None):
super().__init__(config)
self.config = config
Expand All @@ -28,8 +28,8 @@ def __init__(self, config, retriever=None, generator=None):
'triple_save_path': os.path.join(config['save_dir'], 'save_triples.json'),
'triple_load_path': None
}
if 'kg_config' in config and config['kg_config'] is not None:
default_setting.update(config['kg_config'])
if 'trace_config' in config and config['trace_config'] is not None:
default_setting.update(config['trace_config'])
self.kg_setting = default_setting

self.num_examplars = self.kg_setting['num_examplars']
Expand Down
4 changes: 2 additions & 2 deletions flashrag/utils/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -146,10 +146,10 @@ def get_refiner(config, retriever=None, generator=None):
importlib.import_module("flashrag.refiner"),
"SelectiveContextRefiner"
)(config)
elif 'kg' in refiner_name.lower():
elif 'kg-trace' in refiner_name.lower():
return getattr(
importlib.import_module("flashrag.refiner"),
"KGRefiner"
"KGTraceRefiner"
)(config, retriever, generator)
else:
assert False, "No implementation!"
Expand Down

0 comments on commit b9d6fa1

Please sign in to comment.