Parse the dstore path to extract model_type, eval_subset, and dimension. Example: /fs-computility/plm/shared/jqcao/projects/neuralKNN/dstore/Qwen2.5-7B/reviews/dstore_qwen2_train_3584.arrow
(dstore_path)
| 18 | from transformers import Qwen2Model |
| 19 | |
| 20 | def parse_dstore_path(dstore_path): |
| 21 | """ |
| 22 | Parse the dstore path to extract model_type, eval_subset, and dimension. |
| 23 | |
| 24 | Example: |
| 25 | /fs-computility/plm/shared/jqcao/projects/neuralKNN/dstore/Qwen2.5-7B/reviews/dstore_qwen2_train_3584.arrow |
| 26 | """ |
| 27 | # Extract directory and filename |
| 28 | dstore_dir = os.path.dirname(dstore_path) |
| 29 | filename = os.path.basename(dstore_path) |
| 30 | |
| 31 | # Extract dimension from filename (the number before .arrow) |
| 32 | dimension_match = re.search(r'_(\d+)\.arrow$', filename) |
| 33 | if not dimension_match: |
| 34 | raise ValueError(f"Could not extract dimension from filename: {filename}") |
| 35 | dimension = int(dimension_match.group(1)) |
| 36 | |
| 37 | # Extract eval_subset from filename (usually between model and dimension) |
| 38 | # Format typically: dstore_model_subset_dimension.arrow |
| 39 | parts = filename.split('_') |
| 40 | if len(parts) < 3: |
| 41 | raise ValueError(f"Unexpected filename format: {filename}") |
| 42 | eval_subset = parts[-2] # Second to last part (before dimension) |
| 43 | |
| 44 | # Get model_type from the parent directory structure |
| 45 | # The parent directory of the dstore directory is typically the model name |
| 46 | model_dir = os.path.basename(dstore_dir) |
| 47 | parent_dir = os.path.basename(os.path.dirname(dstore_dir)) |
| 48 | model_type = parent_dir |
| 49 | |
| 50 | return { |
| 51 | "dstore_dir": dstore_dir, |
| 52 | "model_type": model_type, |
| 53 | "eval_subset": eval_subset, |
| 54 | "dimension": dimension |
| 55 | } |
| 56 | |
| 57 | def get_index_path(dstore_info): |
| 58 | """Generate the path for the FAISS index file.""" |