Update endpoints.py

endpoints.py

@@ -11,7 +11,8 @@ from routers import inference, training
 from routers.donut_inference import process_document_donut
 from huggingface_hub import login
 import os
-# login(settings.huggingface_key)
+
+# Login using Hugging Face token from environment
 login(os.getenv("HUGGINGFACE_KEY"))
 
 app = FastAPI(openapi_url="/api/v1/sparrow-ml/openapi.json", docs_url="/api/v1/sparrow-ml/docs")
@@ -31,47 +32,44 @@ router = APIRouter()
 
 def count_values(obj):
     if isinstance(obj, dict):
-        count = 0
-        for value in obj.values():
-            count += count_values(value)
-        return count
+        return sum(count_values(v) for v in obj.values())
     elif isinstance(obj, list):
-        count = 0
-        for item in obj:
-            count += count_values(item)
-        return count
+        return sum(count_values(i) for i in obj)
     else:
         return 1
 
 @router.post("/inference")
-async def run_inference(file: Optional[UploadFile] = File(None), image_url: Optional[str] = Form(None),
-                        shipper_id: int = Form(...), model_in_use: str = Form('donut')):
-
+async def run_inference(
+    file: Optional[UploadFile] = File(None), 
+    image_url: Optional[str] = Form(None),
+    shipper_id: int = Form(...), 
+    model_in_use: str = Form('donut')
+):
     result = []
-    model_url = settings.get_model_url(shipper_id)  # Get the correct model URL based on shipper_id
-    
+
+    # Dynamically select model
+    model_url = settings.get_model_url(shipper_id)
+    model_name = model_url.replace("https://huggingface.co/spaces/", "")
+    print(f"[DEBUG] Using model: {model_name}")
+
     if file:
-        # Ensure the uploaded file is a JPG image
         if file.content_type not in ["image/jpeg", "image/jpg"]:
             return {"error": "Invalid file type. Only JPG images are allowed."}
-
+        
         image = Image.open(BytesIO(await file.read()))
-        processing_time = 0
-        if model_in_use == 'donut':
-            result, processing_time = process_document_donut(image, model_url)  # Pass model_url to the function
-        utils.log_stats(settings.inference_stats_file, [processing_time, count_values(result), file.filename, settings.model])
+        result, processing_time = process_document_donut(image, model_url)
+        utils.log_stats(settings.inference_stats_file, [processing_time, count_values(result), file.filename, model_name])
         print(f"Processing time: {processing_time:.2f} seconds")
+    
     elif image_url:
-        # test image url: https://raw.githubusercontent.com/katanaml/sparrow/main/sparrow-data/docs/input/invoices/processed/images/invoice_10.jpg
         with urllib.request.urlopen(image_url) as url:
             image = Image.open(BytesIO(url.read()))
-
-        processing_time = 0
-        if model_in_use == 'donut':
-            result, processing_time = process_document_donut(image, model_url)
+        
+        result, processing_time = process_document_donut(image, model_url)
         file_name = image_url.split("/")[-1]
-        utils.log_stats(settings.inference_stats_file, [processing_time, count_values(result), file_name, settings.model])
+        utils.log_stats(settings.inference_stats_file, [processing_time, count_values(result), file_name, model_name])
         print(f"Processing time inference: {processing_time:.2f} seconds")
+    
     else:
         result = {"info": "No input provided"}
 
@@ -80,8 +78,6 @@ async def run_inference(file: Optional[UploadFile] = File(None), image_url: Opti
 @router.get("/statistics")
 async def get_statistics():
     file_path = settings.inference_stats_file
-
-    # Check if the file exists, and read its content
     if os.path.exists(file_path):
         with open(file_path, 'r') as file:
             try:
@@ -90,7 +86,6 @@ async def get_statistics():
                 content = []
     else:
         content = []
-
     return content
 
 @app.get("/")