Update .gitignore to include .env, add model download script, enhance save_data endpoint for batch processing, and modify model_type formatting in index.html

.gitignore

@@ -1,2 +1,3 @@
 .venv
 __pycache__/
+.env

Dockerfile

@@ -24,4 +24,6 @@ WORKDIR $HOME/app
 # Copy the current directory contents into the container at $HOME/app setting the owner to the user
 COPY --chown=user . $HOME/app
 
+RUN python3 $HOME/app/download_models.py
+
 CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

download_models.py

@@ -0,0 +1,10 @@
+from transformers import AutoModel, AutoTokenizer
+
+model_names = [
+    "WhereIsAI/UAE-Large-V1",
+    "BAAI/bge-large-en-v1.5",
+]
+
+for name in model_names:
+    AutoModel.from_pretrained(name)
+    AutoTokenizer.from_pretrained(name)

main.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, HTTPException, Depends, status
+from fastapi import FastAPI, HTTPException, Depends, status, BackgroundTasks
 from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
@@ -19,6 +19,8 @@ import os
 import logging
 from functools import lru_cache
 from diskcache import Cache
+import json
+import asyncio
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -36,8 +38,8 @@ app = FastAPI()
 cache = Cache('./cache')
 
 # JWT Configuration
-SECRET_KEY = os.environ.get("prime_auth", "c0369f977b69e717dc16f6fc574039eb2b1ebde38014d2be")
-REFRESH_SECRET_KEY = os.environ.get("prolonged_auth", "916018771b29084378c9362c0cd9e631fd4927b8aea07f91")
+SECRET_KEY = os.environ.get("PRIME_AUTH", "c0369f977b69e717dc16f6fc574039eb2b1ebde38014d2be")
+REFRESH_SECRET_KEY = os.environ.get("PROLONGED_AUTH", "916018771b29084378c9362c0cd9e631fd4927b8aea07f91")
 ALGORITHM = "HS256"
 ACCESS_TOKEN_EXPIRE_MINUTES = 30
 REFRESH_TOKEN_EXPIRE_DAYS = 7
@@ -322,88 +324,99 @@ async def search(
             detail=f"Search failed: {str(e)}"
         )
 
+# new constants
+QUEUE_FILE = "./save_queue.jsonl"
+PUSH_INTERVAL_S = 300  # seconds
+QUEUE_THRESHOLD = 1000
+MAX_PUSH_INTERVAL_S = 47 * 3600  # 44 hours
+
+# background task to batch-push queued records
+async def _hf_sync_loop():
+    # authenticate once for private repo access
+    hf_token = os.environ.get("HF_TOKEN")
+    if not hf_token:
+        logging.error("HF_TOKEN not set for Hugging Face authentication")
+        return
+    login(token=hf_token)
+
+    last_push_time = datetime.now(timezone.utc).timestamp()
+
+    while True:
+        await asyncio.sleep(PUSH_INTERVAL_S)
+        try:
+            # Count lines in queue file
+            if not os.path.exists(QUEUE_FILE):
+                continue
+            with open(QUEUE_FILE, "r") as f:
+                lines = f.read().splitlines()
+            queue_len = len(lines)
+            now = datetime.now(timezone.utc).timestamp()
+            time_since_last_push = now - last_push_time
+
+            # print(f"Queue length: {queue_len}, Time since last push: {time_since_last_push}")
+            # Only push if threshold met or max interval 
+            if queue_len >= QUEUE_THRESHOLD or time_since_last_push >= MAX_PUSH_INTERVAL_S:
+                if not lines:
+                    last_push_time = now
+                    continue
+                new_records = [json.loads(l) for l in lines]
+                # load remote dataset with auth
+                dataset = load_dataset(
+                    "HumbleBeeAI/al-ghazali-rag-retrieval-evaluation",
+                    split="train"
+                )
+                data = dataset.to_dict()
+                # append new records
+                for rec in new_records:
+                    for k, v in rec.items():
+                        data.setdefault(k, []).append(v)
+                updated = Dataset.from_dict(data)
+                # push with token
+                updated.push_to_hub(
+                    "HumbleBeeAI/al-ghazali-rag-retrieval-evaluation",
+                    token=hf_token
+                )
+                # clear queue
+                open(QUEUE_FILE, "w").close()
+                last_push_time = now
+        except Exception as e:
+            logging.error(f"Background sync failed: {e}")
+
+# replace existing startup_event
+@app.on_event("startup")
+async def startup_event():
+    os.makedirs("./cache", exist_ok=True)
+    Path(QUEUE_FILE).touch(exist_ok=True)
+    # start background sync loop
+    asyncio.create_task(_hf_sync_loop())
+
+# replace existing /save endpoint
 @app.post("/save")
 async def save_data(
     save_input: SaveBatchInput,
     username: str = Depends(verify_access_token)
 ):
-    try:
-        # Login to Hugging Face
-        hf_token = os.environ.get("al_ghazali_rag_retrieval_evaluation")
-        if not hf_token:
-            raise HTTPException(
-                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-                detail="Hugging Face API token not found"
-            )
-        login(token=hf_token)
-        
-        # Prepare data for saving
-        data = {
-            "user_type": [],
-            "username": [],
-            "query": [],
-            "retrieved_text": [],
-            "model_type": [],
-            "reaction": [],
-            "timestamp": [],
-            "confidence_score": []
-        }
-        
-        # Add each item to the data dict
-        for item in save_input.items:
-            data["user_type"].append(item.user_type)
-            data["username"].append(item.username)
-            data["query"].append(item.query)
-            data["retrieved_text"].append(item.retrieved_text)
-            data["model_type"].append(item.model_type)
-            data["reaction"].append(item.reaction)
-            data["timestamp"].append(datetime.now(timezone.utc).isoformat().replace('+00:00', 'Z'))
-            data["confidence_score"].append(item.confidence_score)
-        
-        try:
-            # Load existing dataset and merge
-            dataset = load_dataset(
-                "HumbleBeeAI/al-ghazali-rag-retrieval-evaluation", 
-                split="train"
-            )
-            existing_data = dataset.to_dict()
-            
-            # Add new data
-            for key in data:
-                if key not in existing_data:
-                    existing_data[key] = [
-                         "" if key in ["timestamp"] else 
-                         0.0 if key in ["confidence_score"] else None
-                     ] * len(next(iter(existing_data.values())))
-                existing_data[key].extend(data[key])
-        
-        except Exception as e:
-            logging.warning(f"Could not load existing dataset, creating new one: {str(e)}")
-            existing_data = data
-        
-        # Create and push dataset
-        updated_dataset = Dataset.from_dict(existing_data)
-        updated_dataset.push_to_hub(
-            "HumbleBeeAI/al-ghazali-rag-retrieval-evaluation"
-        )
-        
-        return {"message": "Data saved successfully"}
-    
-    except Exception as e:
-        logging.error(f"Save error: {str(e)}")
-        raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=f"Failed to save data: {str(e)}"
-        )
+    records = []
+    for item in save_input.items:
+        records.append({
+            "user_type": item.user_type,
+            "username": item.username,
+            "query": item.query,
+            "retrieved_text": item.retrieved_text,
+            "model_type": item.model_type,
+            "reaction": item.reaction,
+            "timestamp": datetime.now(timezone.utc).isoformat().replace('+00:00','Z'),
+            "confidence_score": item.confidence_score
+        })
+    # append to local queue
+    with open(QUEUE_FILE, "a") as f:
+        for r in records:
+            f.write(json.dumps(r) + "\n")
+    return {"message": "Your data is queued for batch upload."}
 
 # Make sure to keep the static files mounting
 app.mount("/home", StaticFiles(directory="static", html=True), name="home")
 
-# Startup event to create cache directory if it doesn't exist
-@app.on_event("startup")
-async def startup_event():
-    os.makedirs("./cache", exist_ok=True)
-
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)

static/index.html

@@ -219,7 +219,7 @@ class LoginResponse {
                                     "username": "user1332",
                                     "query": "What is the seventh test of lovers of God?",
                                     "retrieved_text": "The seventh test is that lovers of God will love those who obey Him and hate the infidels and the disobedient...",
-                                    "model_type": "WhereIsAI/UAE-Large-V1",
+                                    "model_type": "WhereIsAI_UAE_Large_V1",
                                     "reaction": "positive",
                                     "confidence_score": 0.95
                                 },
@@ -228,7 +228,7 @@ class LoginResponse {
                                     "username": "user1332",
                                     "query": "What is the seventh test of lovers of God?",
                                     "retrieved_text": "The seventh test is that lovers of God will love those who obey Him and hate the infidels and the disobedient...",
-                                    "model_type": "BAAI/bge-large-en-v1.5",
+                                    "model_type": "BAAI_bge_large_en_v1.5",
                                     "reaction": "positive",
                                     "confidence_score": 0.92
                                 }