Update app.py

app.py

@@ -1,8 +1,8 @@
+import os
 import json
+import re
 import gradio as gr
 import requests
-import os
-import re
 
 # Hugging Face API details
 API_URL = "https://api-inference.huggingface.co/models/numind/NuExtract-1.5"
@@ -10,39 +10,18 @@ api_token = os.environ.get("HF_TOKEN", "")  # Get token from environment variabl
 headers = {"Authorization": f"Bearer {api_token}"}
 
 
-# Test API Connection
-def test_api_connection():
-    try:
-        # Test with a simple GET request
-        response = requests.get("https://api-inference.huggingface.co/models/numind/NuExtract-1.5")
-
-        # Check if the connection was successful
-        if response.status_code == 200:
-            print("✅ Connection to Hugging Face API successful!")
-        else:
-            print(f"⚠️ API returned status code {response.status_code}: {response.text}")
-    except requests.exceptions.RequestException as e:
-        print(f"❌ Connection failed: {str(e)}")
-
-
-# Make the API request
 def query_api(payload):
     try:
         response = requests.post(API_URL, headers=headers, json=payload)
-        # Debug logs
         print("API STATUS CODE:", response.status_code)
         print("RAW RESPONSE:", response.text)
-
         return response.json()
     except Exception as e:
-        print("Error during API call:", e)
-        return {"error": f"Could not decode JSON: {str(e)}"}
+        return {"error": f"Request failed: {str(e)}"}
 
 
-# Extract structure from the template and text
 def extract_structure(template, text):
     try:
-        # Format the input for NuExtract
         prompt = f"<|input|>\n### Template:\n{template}\n### Text:\n{text}\n\n<|output|>"
 
         payload = {
@@ -56,24 +35,19 @@ def extract_structure(template, text):
 
         response = query_api(payload)
 
-        # Check for API error
         if isinstance(response, dict) and "error" in response:
-            return f"API Error: {response['error']}", "{}", "<p>Error occurred: {response['error']}</p>"
+            return f"API Error: {response['error']}", "{}", f"<p>Error: {response['error']}</p>"
 
-        # Get generated text
         if isinstance(response, list) and len(response) > 0:
             output = response[0].get("generated_text", "")
-            print("Generated Text:", output)  # Optional debugging
+            print("Generated Text:", output)
 
-            # Try to extract after <|output|>
             if "<|output|>" in output:
                 result = output.split("<|output|>")[-1].strip()
             else:
-                # Try to extract JSON-like structure using regex
-                json_match = re.search(r'({[\s\S]+})', output)
-                result = json_match.group(1) if json_match else output.strip()
+                match = re.search(r'({[\s\S]+})', output)
+                result = match.group(1) if match else output.strip()
 
-            # Attempt to format JSON nicely
             try:
                 parsed = json.loads(result)
                 result = json.dumps(parsed, indent=2)
@@ -82,22 +56,19 @@ def extract_structure(template, text):
 
             highlighted = f"<p>✅ Successfully processed input of length {len(text)} characters.</p>"
             return "✅ Extraction Complete", result, highlighted
-        else:
-            return "⚠️ Unexpected API Response", json.dumps(response, indent=2), "<p>Please check the API response format.</p>"
+
+        return "⚠️ Unexpected API Response", json.dumps(response, indent=2), "<p>Unexpected format.</p>"
 
     except Exception as e:
         return f"❌ Error: {str(e)}", "{}", f"<p>Processing failed: {str(e)}</p>"
 
 
-# Gradio Interface
+# Gradio App
 with gr.Blocks() as demo:
     gr.Markdown("# 🧠 NuExtract-1.5 Information Extractor")
 
     if not api_token:
-        gr.Markdown("## ⚠️ No API token found. Set `HF_TOKEN` in the Space secrets.")
-
-    # Call test connection before launching the Gradio interface
-    test_api_connection()
+        gr.Markdown("## ⚠️ No API token found. Please set `HF_TOKEN` in environment variables.")
 
     with gr.Row():
         with gr.Column():
@@ -148,6 +119,26 @@ with gr.Blocks() as demo:
         [template_input, text_input]
     )
 
-if __name__ == "__main__":
-    demo.launch()
 
+def test_api_connection():
+    print("===== Application Startup =====")
+    if not api_token:
+        print("❌ HF_TOKEN not set. Please set your API token.")
+    else:
+        test_payload = {
+            "inputs": "<|input|>\n### Template:\n{\"test\": \"\"}\n### Text:\nHello world\n\n<|output|>",
+            "parameters": {
+                "max_new_tokens": 100,
+                "temperature": 0.01
+            }
+        }
+        response = query_api(test_payload)
+        if isinstance(response, list):
+            print("✅ Connection to Hugging Face API successful!")
+        else:
+            print("⚠️ API may not be returning expected format:", response)
+
+
+if __name__ == "__main__":
+    test_api_connection()
+    demo.launch(debug=True)  # You can add share=True or server_name/port if needed