64 lines
2.3 KiB
Python
64 lines
2.3 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import csv
|
||
|
|
|
||
|
|
from app.domain.priorities import ORIGIN_TO_PRIORITY, PRIORITY_LOW_CONFIDENCE
|
||
|
|
from app.settings import load_settings
|
||
|
|
from app.storage.db import connect
|
||
|
|
from app.storage.repositories import Repository
|
||
|
|
|
||
|
|
|
||
|
|
def flags_for_origin(origin: str) -> tuple[bool, bool, bool, bool]:
|
||
|
|
is_bounty = origin in {"manual_bounty", "bounty"}
|
||
|
|
is_promote = origin in {"manual_promote", "promote"}
|
||
|
|
downloaded = is_bounty or origin in {"manual_downloaded_by_others", "downloaded_by_others"}
|
||
|
|
self_download = origin in {"hf_seed", "self_download"}
|
||
|
|
return is_bounty, is_promote, downloaded, self_download
|
||
|
|
|
||
|
|
|
||
|
|
def load_seed(path: str, origin: str, repo: Repository) -> int:
|
||
|
|
priority = ORIGIN_TO_PRIORITY.get(origin, PRIORITY_LOW_CONFIDENCE)
|
||
|
|
count = 0
|
||
|
|
with open(path, newline="", encoding="utf-8") as fh:
|
||
|
|
reader = csv.DictReader(fh)
|
||
|
|
for row in reader:
|
||
|
|
model_id = (row.get("model_id") or "").strip()
|
||
|
|
if not model_id:
|
||
|
|
continue
|
||
|
|
is_bounty, is_promote, downloaded, self_download = flags_for_origin(origin)
|
||
|
|
repo.upsert_candidate(
|
||
|
|
model_id=model_id,
|
||
|
|
source=(row.get("source") or "HUGGING_FACE").strip(),
|
||
|
|
origin=origin,
|
||
|
|
priority=priority,
|
||
|
|
is_bounty=is_bounty,
|
||
|
|
is_promote=is_promote,
|
||
|
|
known_downloaded_by_others=downloaded,
|
||
|
|
self_download_allowed=self_download,
|
||
|
|
max_model_len=int(row.get("max_model_len") or 1024),
|
||
|
|
target_gpu_aliases=(row.get("gpu_aliases") or "").strip() or None,
|
||
|
|
notes=row.get("notes"),
|
||
|
|
)
|
||
|
|
count += 1
|
||
|
|
return count
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--file", required=True)
|
||
|
|
parser.add_argument("--origin", required=True, choices=sorted(ORIGIN_TO_PRIORITY))
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
settings = load_settings(require_secrets=False)
|
||
|
|
conn = connect(settings.db_path)
|
||
|
|
try:
|
||
|
|
count = load_seed(args.file, args.origin, Repository(conn))
|
||
|
|
print(f"imported {count} candidates from {args.file} as {args.origin}")
|
||
|
|
finally:
|
||
|
|
conn.close()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|