Update app.py

app.py

@@ -2,93 +2,108 @@ import json
 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"
 api_token = os.environ.get("HF_TOKEN", "")  # Get token from environment variable
-
 headers = {"Authorization": f"Bearer {api_token}"}
 
+
 def query_api(payload):
     response = requests.post(API_URL, headers=headers, json=payload)
-    return response.json()
+
+    # Debug logs
+    print("API STATUS CODE:", response.status_code)
+    print("RAW RESPONSE:", response.text)
+
+    try:
+        return response.json()
+    except Exception as e:
+        return {"error": f"Could not decode JSON: {str(e)}"}
+
 
 def extract_structure(template, text):
     try:
-        # Format the input following NuExtract's format
+        # Format the input for NuExtract
         prompt = f"<|input|>\n### Template:\n{template}\n### Text:\n{text}\n\n<|output|>"
-        
-        # Call the API
+
         payload = {
             "inputs": prompt,
             "parameters": {
                 "max_new_tokens": 2000,
-                "temperature": 0.01,  # Nearly deterministic as recommended
+                "temperature": 0.01,
                 "return_full_text": True
             }
         }
-        
+
         response = query_api(payload)
-        
-        # Check for errors
+
+        # Check for API error
         if isinstance(response, dict) and "error" in response:
             return f"API Error: {response['error']}", "{}", f"<p>Error occurred: {response['error']}</p>"
-            
-        # Extract result - the API returns the full text so we need to split it
+
+        # Get generated text
         if isinstance(response, list) and len(response) > 0:
             output = response[0].get("generated_text", "")
-            result = output.split("<|output|>")[1] if "<|output|>" in output else output
-            
-            # Try to parse as JSON to format it nicely
+            print("Generated Text:", output)  # Optional debugging
+
+            # 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()
+
+            # Attempt to format JSON nicely
             try:
                 parsed = json.loads(result)
                 result = json.dumps(parsed, indent=2)
-            except:
+            except Exception:
                 pass
-                
-            # Create a simple highlight
-            highlighted = f"<p>Successfully processed text of length {len(text)} characters</p>"
-            
-            return "Processing complete", result, highlighted
+
+            highlighted = f"<p>✅ Successfully processed input of length {len(text)} characters.</p>"
+            return "✅ Extraction Complete", result, highlighted
         else:
-            return "Unexpected API response", str(response), "<p>Please check API token and try again</p>"
-            
+            return "⚠️ Unexpected API Response", json.dumps(response, indent=2), "<p>Please check the API response format.</p>"
+
     except Exception as e:
-        return f"Error: {str(e)}", "{}", f"<p>Processing failed: {str(e)}</p>"
+        return f"❌ Error: {str(e)}", "{}", f"<p>Processing failed: {str(e)}</p>"
 
-# Create interface
+
+# Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# NuExtract-1.5 Demo")
-    
+    gr.Markdown("# 🧠 NuExtract-1.5 Information Extractor")
+
     if not api_token:
-        gr.Markdown("## ⚠️ No API token found. Set HF_TOKEN in Space secrets.")
-    
+        gr.Markdown("## ⚠️ No API token found. Set `HF_TOKEN` in the Space secrets.")
+
     with gr.Row():
         with gr.Column():
             template_input = gr.Textbox(
-                label="Template (JSON)", 
+                label="Template (JSON)",
                 value='{"name": "", "email": ""}',
                 lines=5
             )
             text_input = gr.Textbox(
-                label="Input Text", 
+                label="Input Text",
                 value="Contact: John Smith ([email protected])",
                 lines=10
             )
             submit_btn = gr.Button("Extract Information")
-        
+
         with gr.Column():
             progress_output = gr.Textbox(label="Progress")
             result_output = gr.Textbox(label="Extracted Information")
-            html_output = gr.HTML(label="Highlighted Text")
-    
+            html_output = gr.HTML(label="Info")
+
     submit_btn.click(
         fn=extract_structure,
         inputs=[template_input, text_input],
         outputs=[progress_output, result_output, html_output]
     )
 
-    # Examples
     gr.Examples(
         [
             [
@@ -114,4 +129,4 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()