Update app.py

app.py

@@ -33,9 +33,16 @@ if not hf_token:
     logger.error("HF_TOKEN environment variable not set")
     raise HTTPException(status_code=500, detail="HF_TOKEN environment variable not set")
 
-# Initialize Hugging Face Inference Client
-client = InferenceClient(model="Qwen/Qwen2-7B-Instruct", token=hf_token)
-logger.info("Hugging Face Inference Client initialized for Qwen2-7B-Instruct")
+# Initialize Hugging Face Inference Client with primary and fallback models
+primary_model = "Qwen/Qwen2-7B-Instruct"
+fallback_model = "mistral/Mixtral-8x7B-Instruct-v0.1"
+try:
+    client = InferenceClient(model=primary_model, token=hf_token, provider="auto")
+    logger.info(f"Hugging Face Inference Client initialized for {primary_model} with provider='auto'")
+except Exception as e:
+    logger.warning(f"Failed to initialize client for {primary_model}: {str(e)}. Falling back to {fallback_model}")
+    client = InferenceClient(model=fallback_model, token=hf_token, provider="auto")
+    logger.info(f"Hugging Face Inference Client initialized for {fallback_model} with provider='auto'")
 
 # In-memory caches (1-hour TTL)
 raw_text_cache = cachetools.TTLCache(maxsize=100, ttl=3600)
@@ -90,7 +97,7 @@ async def process_pdf_page(img, page_idx):
         return ""
 
 async def process_with_qwen(filename: str, raw_text: str):
-    """Process raw text with Qwen2-7B-Instruct via Hugging Face Inference API."""
+    """Process raw text with Qwen2-7B-Instruct or fallback via Hugging Face Inference API."""
     start_time = time.time()
     logger.info(f"Starting Qwen API processing for {filename}, {log_memory_usage()}")
 
@@ -107,7 +114,7 @@ async def process_with_qwen(filename: str, raw_text: str):
 
     try:
         prompt = f"""
-        Extract key invoice fields as JSON from the raw text. Support English. Detect currency (e.g., USD, INR). Output only the JSON object.
+        Extract key invoice fields as JSON from the raw text. Support English. Detect currency (e.g., USD, INR). Output only valid JSON, with no additional text, comments, or markdown.
         Raw text: {raw_text}
         Output JSON:
         {{
@@ -133,7 +140,7 @@ async def process_with_qwen(filename: str, raw_text: str):
         json_start = llm_output.find("{")
         json_end = llm_output.rfind("}") + 1
         if json_start == -1 or json_end == -1:
-            raise ValueError("No valid JSON found in Qwen API output")
+            raise ValueError("No valid JSON found in API output")
         json_str = llm_output[json_start:json_end]
         structured_data = json.loads(json_str)
         structured_data_cache[text_hash] = structured_data