moar ram

app.py

@@ -1,30 +1,30 @@
 import gradio as gr
 from huggingface_hub import hf_hub_download
-import json
+import json_stream as json
 import gzip
 import urllib
-
-usernames = {}
-
-filepath = hf_hub_download(repo_id="bigcode/the-stack-username-to-repo", filename="username_to_repo.json.gz", repo_type="dataset", revision="v2.0.1")
-with gzip.open(filepath, 'r') as f:
-    usernames["v2.0.1"] = json.loads(f.read().decode('utf-8'))
-
-filepath = hf_hub_download(repo_id="bigcode/the-stack-username-to-repo", filename="username_to_repo.json.gz", repo_type="dataset", revision="v2.0")
-with gzip.open(filepath, 'r') as f:
-    usernames["v2.0"] = json.loads(f.read().decode('utf-8'))
-
-filepath = hf_hub_download(repo_id="bigcode/the-stack-username-to-repo", filename="username_to_repo.json.gz", repo_type="dataset", revision="v1.2")
-with gzip.open(filepath, 'r') as f:
-    usernames["v1.2"] = json.loads(f.read().decode('utf-8'))
-
-filepath = hf_hub_download(repo_id="bigcode/the-stack-username-to-repo", filename="username_to_repo.json.gz", repo_type="dataset", revision="v1.1")
-with gzip.open(filepath, 'r') as f:
-    usernames["v1.1"] = json.loads(f.read().decode('utf-8'))
-
-filepath = hf_hub_download(repo_id="bigcode/the-stack-username-to-repo", filename="username_to_repo.json.gz", repo_type="dataset")
-with gzip.open(filepath, 'r') as f:
-    usernames["v1.0"] = json.loads(f.read().decode('utf-8'))
+from collections import defaultdict
+import gc
+import sys
+
+usernames = defaultdict(dict)
+
+versions = ["v1.0", "v1.1", "v1.2", "v2.0", "v2.0.1"]
+versions = [sys.intern(version) for version in versions]
+
+for version in versions:
+    print(f"Loading {version}")
+    branch = version if version != "v1.0" else "main"
+    filepath = hf_hub_download(repo_id="bigcode/the-stack-username-to-repo", filename="username_to_repo.json.gz", repo_type="dataset", revision=branch)
+    with gzip.open(filepath, 'r') as f:
+        data = json.load(f)
+        for username, repos in data.items():
+            for repo in repos:
+                if repo not in usernames[username]:
+                    usernames[username][repo] = []
+                usernames[username][repo].append(version)
+        del data
+    gc.collect()
 
 text = """\
 ![](https://huggingface.co/spaces/lvwerra/in-the-stack-gr/resolve/main/banner.png)
@@ -77,8 +77,11 @@ def issue_url(username, repos):
 
 def check_username(username, version):
     output_md = ""
-    if username in usernames[version] and len(usernames[version][username])>0:
-        repos = usernames[version][username]
+    repos = []
+    if username in usernames:
+        repos = [repo for repo, versions in usernames[username].items() if version in versions]
+
+    if repos:
         repo_word = "repository" if len(repos)==1 else "repositories"
         if version[:2] == "v2":
             output_md += f"**Yes**, there is code from **{len(repos)} {repo_word}** in The Stack. Check the links to see when it was archived by Software Heritage:\n\n"

requirements.txt

@@ -0,0 +1 @@
+json-stream