Update app.py

app.py

@@ -1,49 +1,67 @@
 import json
-import torch
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import requests
+import os
 
-# Simplified extraction function
-def extract_structure(template, text, progress=None):
+# 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()
+
+def extract_structure(template, text):
     try:
-        # Format the input
+        # Format the input following NuExtract's format
         prompt = f"<|input|>\n### Template:\n{template}\n### Text:\n{text}\n\n<|output|>"
         
-        # Generate prediction
-        input_ids = tokenizer(prompt, return_tensors="pt", truncation=True).to(model.device)
-        output = tokenizer.decode(model.generate(**input_ids, max_new_tokens=2000)[0], skip_special_tokens=True)
-        
-        # Extract result
-        result = output.split("<|output|>")[1]
+        # Call the API
+        payload = {
+            "inputs": prompt,
+            "parameters": {
+                "max_new_tokens": 2000,
+                "temperature": 0.01,  # Nearly deterministic as recommended
+                "return_full_text": True
+            }
+        }
         
-        # Highlight found items in text (simplified)
-        highlighted = f"<p>Processed text of length {len(text)} characters</p>"
+        response = query_api(payload)
         
-        return "Processing complete", result, highlighted
+        # Check for errors
+        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
+        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
+            try:
+                parsed = json.loads(result)
+                result = json.dumps(parsed, indent=2)
+            except:
+                pass
+                
+            # Create a simple highlight
+            highlighted = f"<p>Successfully processed text of length {len(text)} characters</p>"
+            
+            return "Processing complete", result, highlighted
+        else:
+            return "Unexpected API response", str(response), "<p>Please check API token and try again</p>"
+            
     except Exception as e:
-        return f"Error: {str(e)}", "{}", "<p>Processing failed</p>"
-
-# Load model
-model_name = "numind/NuExtract-1.5"
-try:
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name, 
-        torch_dtype=torch.float16,  # Using float16 instead of bfloat16 for better compatibility
-        trust_remote_code=True,
-        device_map="auto"
-    )
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model_loaded = True
-except Exception as e:
-    print(f"Model loading error: {e}")
-    model_loaded = False
+        return f"Error: {str(e)}", "{}", f"<p>Processing failed: {str(e)}</p>"
 
 # Create interface
 with gr.Blocks() as demo:
     gr.Markdown("# NuExtract-1.5 Demo")
     
-    if not model_loaded:
-        gr.Markdown("## ⚠️ Model failed to load. Using dummy mode.")
+    if not api_token:
+        gr.Markdown("## ⚠️ No API token found. Set HF_TOKEN in Space secrets.")
     
     with gr.Row():
         with gr.Column():
@@ -70,12 +88,26 @@ with gr.Blocks() as demo:
         outputs=[progress_output, result_output, html_output]
     )
 
-    # Simple example
+    # Examples
     gr.Examples(
         [
             [
                 '{"name": "", "email": ""}',
                 'Contact: John Smith ([email protected])'
+            ],
+            [
+                '''{
+    "Model": {
+        "Name": "",
+        "Number of parameters": "",
+        "Architecture": []
+    },
+    "Usage": {
+        "Use case": [],
+        "License": ""
+    }
+}''',
+                '''We introduce Mistral 7B, a 7-billion-parameter language model engineered for superior performance and efficiency. Mistral 7B outperforms the best open 13B model (Llama 2) across all evaluated benchmarks, and the best released 34B model (Llama 1) in reasoning, mathematics, and code generation. Our model is released under the Apache 2.0 license.'''
             ]
         ],
         [template_input, text_input]