Update app.py

app.py

@@ -1,57 +1,52 @@
 import gradio as gr
-import torch
 import json
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import requests
+import os
 
-# Initialize model with error handling
-try:
-    tokenizer = AutoTokenizer.from_pretrained("numind/NuExtract-1.5")
-    model = AutoModelForCausalLM.from_pretrained(
-        "numind/NuExtract-1.5",
-        device_map="auto",
-        torch_dtype=torch.float16,
-        trust_remote_code=True
-    )
-    MODEL_LOADED = True
-    print("Model loaded successfully!")
-except Exception as e:
-    MODEL_LOADED = False
-    print(f"Model loading failed: {e}")
+# Use the Hugging Face Inference API instead of loading the model
+API_URL = "https://api-inference.huggingface.co/models/numind/NuExtract-1.5"
+headers = {"Authorization": f"Bearer {os.environ.get('HF_TOKEN', '')}"}
 
 def test_function(template, text):
     print(f"Test function called with template: {template[:30]} and text: {text[:30]}")
     return "Button clicked successfully", "Function was called"
 
 def extract_info(template, text):
-    if not MODEL_LOADED:
-        return "❌ Model not loaded", "{}"
-    
     try:
         # Format prompt according to NuExtract-1.5 requirements
         prompt = f"<|input|>\n### Template:\n{template}\n### Text:\n{text}\n\n<|output|>"
         print(f"Processing with prompt: {prompt[:100]}...")
         
-        # Tokenize
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        # Call API instead of using local model
+        payload = {
+            "inputs": prompt,
+            "parameters": {
+                "max_new_tokens": 1000,
+                "do_sample": False
+            }
+        }
+        
+        print("Calling API...")
+        response = requests.post(API_URL, headers=headers, json=payload)
         
-        # Generate with cache disabled
-        print("Generating output...")
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=1000,
-            do_sample=False,
-            use_cache=False  # This disables the problematic cache
-        )
+        if response.status_code != 200:
+            print(f"API error: {response.status_code}, {response.text}")
+            return f"❌ API Error: {response.status_code}", response.text
         
-        # Decode and extract result
-        print("Decoding output...")
-        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Process result
+        result = response.json()
+        
+        # Handle different response formats
+        if isinstance(result, list) and len(result) > 0:
+            result_text = result[0].get("generated_text", "")
+        else:
+            result_text = str(result)
         
-        # Split at output marker
-        if "<|output|>" in result:
-            json_text = result.split("<|output|>")[1].strip()
+        # Split at output marker if present
+        if "<|output|>" in result_text:
+            json_text = result_text.split("<|output|>")[1].strip()
         else:
-            json_text = result
+            json_text = result_text
         
         # Try to parse as JSON
         print("Parsing JSON...")