Add language parsing functionality and update dependencies

main.py

@@ -15,6 +15,7 @@ from starlette.responses import RedirectResponse
 from cashews import cache
 from datetime import timedelta
 import logging
+from iso639 import Lang
 
 cache.setup("mem://")
 
@@ -93,6 +94,19 @@ async def get_dataset_info(hub_id: str, config: str | None = None):
     return resp.json()
 
 
+@cache(ttl=timedelta(minutes=5))
+async def fetch_rows(url: str) -> list[dict]:
+    response = await async_client.get(url)
+    if response.status_code == 200:
+        data = response.json()
+        return data.get("rows")
+    else:
+        print(f"Failed to fetch data: {response.status_code}")
+        print(url)
+        return []
+
+
+# Function to get random rows from the dataset
 async def get_random_rows(
     hub_id: str,
     total_length: int,
@@ -110,15 +124,8 @@ async def get_random_rows(
         offset = random.randint(0, total_length - rows_per_call)
         url = f"https://datasets-server.huggingface.co/rows?dataset={hub_id}&config={config}&split={split}&offset={offset}&length={rows_per_call}"
         logger.info(f"Fetching {url}")
-        print(url)
-        response = await async_client.get(url)
-        if response.status_code == 200:
-            data = response.json()
-            batch_rows = data.get("rows")
-            rows.extend(batch_rows)
-        else:
-            print(f"Failed to fetch data: {response.status_code}")
-            print(url)
+        batch_rows = await fetch_rows(url)
+        rows.extend(batch_rows)
         if len(rows) >= number_of_rows:
             break
     return [row.get("row") for row in rows]
@@ -181,6 +188,17 @@ def filter_by_frequency(counts_dict: dict, threshold_percent: float = 0.2):
     return {k for k, v in counts_dict.items() if v >= threshold}
 
 
+def try_parse_language(lang: str) -> str | None:
+    try:
+        split = lang.split("_")
+        lang = split[0]
+        lang = Lang(lang)
+        return lang.pt1
+    except Exception as e:
+        logger.error(f"Failed to parse language {lang}: {e}")
+        return None
+
+
 def predict_rows(
     rows, target_column, language_threshold_percent=0.2, return_raw_predictions=False
 ):
@@ -196,8 +214,13 @@ def predict_rows(
         langues_counts, threshold_percent=language_threshold_percent
     )
     filtered_dict = {k: v for k, v in predictions_by_lang.items() if k in keys_to_keep}
+    raw_model_prediction_summary = dict(valmap(get_mean_score, filtered_dict))
+    parsed_langs = {
+        try_parse_language(k): v for k, v in raw_model_prediction_summary.items()
+    }
     default_data = {
-        "predictions": dict(valmap(get_mean_score, filtered_dict)),
+        "language_prediction_summary": parsed_langs,
+        "raw_model_prediction_summary": raw_model_prediction_summary,
         "hub_id": "hub_id",
         "config": "config",
     }

requirements.in

@@ -8,4 +8,5 @@ huggingface_hub
 python-dotenv
 rich
 toolz
-uvicorn[standard]
+uvicorn[standard]
+iso639-lang

requirements.txt

@@ -51,6 +51,8 @@ idna==3.6
     #   anyio
     #   httpx
     #   requests
+iso639-lang==2.2.2
+    # via -r requirements.in
 markdown-it-py==3.0.0
     # via rich
 mdurl==0.1.2