Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 import json
-import requests
 import os
 import pandas as pd
 import folium
@@ -15,6 +14,8 @@ import tempfile
 import warnings
 import string
 import spaces
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 warnings.filterwarnings("ignore")
 
@@ -26,9 +27,16 @@ MAP_TILES = {
     }
 }
 
-# NuExtract API configuration
-API_URL = "https://api-inference.huggingface.co/models/numind/NuExtract-1.5-tiny"
-headers = {"Authorization": f"Bearer {os.environ.get('HF_TOKEN', '')}"}
+# Model configuration
+MODEL_NAME = "numind/NuExtract-1.5-tiny"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+TORCH_DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32
+MAX_INPUT_LENGTH = 20000  # For sliding window processing
+MAX_NEW_TOKENS = 1000
+
+# Global model variables
+tokenizer = None
+model = None
 
 class SafeGeocoder:
     def __init__(self):
@@ -42,7 +50,7 @@ class SafeGeocoder:
         elapsed = current_time - self.last_request
         if elapsed < 1.0:
             time.sleep(1.0 - elapsed)
-        self.last_request = time.time()
+        self.last_request = current_time
     
     def get_coords(self, location: str):
         if not location or pd.isna(location):
@@ -67,82 +75,59 @@ class SafeGeocoder:
             self.cache[location] = None
             return None
 
-# Function to just load the model
-
-
 def load_model():
+    global tokenizer, model
     try:
         # Generate a random location and text each time
         random_city = random.choice(["Berlin", "Paris", "London", "Tokyo", "Rome", "Madrid"])
         random_suffix = ''.join(random.choices(string.ascii_lowercase, k=5))
         test_text = f"Test in {random_city}_{random_suffix}."
-        
         test_template = '{"test_location": ""}'
         
-        prompt = f"<|input|>\n### Template:\n{test_template}\n### Text:\n{test_text}\n\n<|output|>"
-        
-        # Send request with randomized input
-        payload = {
-            "inputs": prompt,
-            "parameters": {
-                "max_new_tokens": 50,
-                "do_sample": False
-            }
-        }
-        
-        response = requests.post(API_URL, headers=headers, json=payload)
-        
-        if response.status_code == 503:
-            response_json = response.json()
-            if "error" in response_json and "loading" in response_json["error"]:
-                estimated_time = response_json.get("estimated_time", "unknown")
-                return f"⏳ Modell lädt... (ca. {int(float(estimated_time)) if isinstance(estimated_time, (int, float, str)) else 'unbekannt'} Sekunden)"
+        # Initialize model if not already loaded
+        if model is None:
+            tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL_NAME,
+                torch_dtype=TORCH_DTYPE,
+                trust_remote_code=True,
+                device_map="auto"
+            ).eval()
+            print(f"✅ Loaded {MODEL_NAME} on {DEVICE}")
         
-        if response.status_code == 200:
-            result = response.json()
-            
-            if isinstance(result, list) and len(result) > 0:
-                result_text = result[0].get("generated_text", "")
-                
-                # Check if response contains the random city we included
-                if "<|output|>" in result_text and random_city in result_text:
-                    return "✅ Modell erfolgreich geladen und getestet! Sie können jetzt mit der Extraktion beginnen."
+        # Test the model
+        prompt = f"<|input|>\n### Template:\n{test_template}\n### Text:\n{test_text}\n\n<|output|>"
+        inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+        outputs = model.generate(**inputs, max_new_tokens=50)
+        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        return "⚠️ Modell-Test nicht erfolgreich. Bitte versuchen Sie es erneut in einigen Sekunden."
+        if "<|output|>" in result and random_city in result:
+            return "✅ Modell erfolgreich geladen und getestet! Sie können jetzt mit der Extraktion beginnen."
+        return "⚠️ Modell-Test nicht erfolgreich. Bitte versuchen Sie es erneut."
             
     except Exception as e:
         return f"❌ Fehler beim Laden des Modells: {str(e)}"
+
 @spaces.GPU
-        
 def extract_info(template, text):
+    global tokenizer, model
+    
+    if model is None:
+        return "❌ Modell nicht geladen", "Bitte zuerst das Modell laden (1. Schritt)"
+    
     try:
         prompt = f"<|input|>\n### Template:\n{template}\n### Text:\n{text}\n\n<|output|>"
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=MAX_INPUT_LENGTH).to(DEVICE)
         
-        payload = {
-            "inputs": prompt,
-            "parameters": {
-                "max_new_tokens": 1000,
-                "do_sample": False
-            }
-        }
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=MAX_NEW_TOKENS,
+            temperature=0.0,
+            do_sample=False,
+            pad_token_id=tokenizer.eos_token_id
+        )
         
-        response = requests.post(API_URL, headers=headers, json=payload)
-        
-        if response.status_code == 503:
-            response_json = response.json()
-            if "error" in response_json and "loading" in response_json["error"]:
-                estimated_time = response_json.get("estimated_time", "unknown")
-                return f"⏳ Modell lädt... (ca. {int(float(estimated_time)) if isinstance(estimated_time, (int, float, str)) else 'unbekannt'} Sekunden)", "Bitte versuchen Sie es in einigen Minuten erneut oder nutzen Sie den 'Modell laden' Button"
-        
-        if response.status_code != 200:
-            return f"❌ API Fehler: {response.status_code}", response.text
-        
-        result = response.json()
-        
-        if isinstance(result, list) and len(result) > 0:
-            result_text = result[0].get("generated_text", "")
-        else:
-            result_text = str(result)
+        result_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
         if "<|output|>" in result_text:
             json_text = result_text.split("<|output|>")[1].strip()
@@ -152,10 +137,10 @@ def extract_info(template, text):
         try:
             extracted = json.loads(json_text)
             formatted = json.dumps(extracted, indent=2)
+            return "✅ Erfolgreich extrahiert", formatted
         except json.JSONDecodeError:
             return "❌ JSON Parsing Fehler", json_text
             
-        return "✅ Erfolgreich extrahiert", formatted
     except Exception as e:
         return f"❌ Fehler: {str(e)}", "{}"