Update app.py

app.py

@@ -16,7 +16,6 @@ import string
 import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers import AutoConfig
 import torch
 
@@ -84,87 +83,73 @@ class SafeGeocoder:
 def load_model():
     global tokenizer, model
     try:
-        # First ensure we have the right tokenizer class available
-        from transformers import Qwen2Tokenizer
-    except ImportError:
-        # Fallback to AutoTokenizer if specific import fails
-        pass
-        
-    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": ""}'
-        
-        # Initialize model if not already loaded
         if model is None:
-            # Load config first to check for tokenizer class
-            config = AutoConfig.from_pretrained(MODEL_NAME, trust_remote_code=True)
-            
-            # Load tokenizer with explicit class if needed
-            if hasattr(config, "tokenizer_class"):
-                tokenizer = AutoTokenizer.from_pretrained(
-                    MODEL_NAME,
-                    trust_remote_code=True,
-                    tokenizer_class=config.tokenizer_class
-                )
-            else:
-                tokenizer = AutoTokenizer.from_pretrained(
-                    MODEL_NAME,
-                    trust_remote_code=True
-                )
+            # Special handling for NuExtract tokenizer
+            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"
+                device_map="auto",
+                trust_remote_code=True
             ).eval()
+            
             print(f"✅ Loaded {MODEL_NAME} on {DEVICE}")
-        
-        # 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)
-        
-        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."
             
+            # Test the model
+            test_text = "Test in Berlin."
+            test_template = '{"test_location": ""}'
+            prompt = f"<|input|>\n### Template:\n{test_template}\n### Text:\n{test_text}\n\n<|output|>"
+            
+            inputs = tokenizer(prompt, return_tensors="pt", max_length=20000, truncation=True).to(DEVICE)
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=50,
+                temperature=0.0,
+                do_sample=False
+            )
+            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            if "<|output|>" in result and "Berlin" in result:
+                return "✅ Modell erfolgreich geladen und getestet!"
+            
+        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)"
+        return "❌ Modell nicht geladen", "Bitte zuerst das Modell laden"
     
     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)
+        inputs = tokenizer(
+            prompt, 
+            return_tensors="pt", 
+            truncation=True, 
+            max_length=20000
+        ).to(DEVICE)
         
         outputs = model.generate(
             **inputs,
-            max_new_tokens=MAX_NEW_TOKENS,
+            max_new_tokens=1000,
             temperature=0.0,
             do_sample=False,
             pad_token_id=tokenizer.eos_token_id
         )
         
         result_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        if "<|output|>" in result_text:
-            json_text = result_text.split("<|output|>")[1].strip()
-        else:
-            json_text = result_text
+        json_text = result_text.split("<|output|>")[1].strip() if "<|output|>" in result_text else result_text
         
         try:
             extracted = json.loads(json_text)
-            formatted = json.dumps(extracted, indent=2)
-            return "✅ Erfolgreich extrahiert", formatted
+            return "✅ Erfolgreich extrahiert", json.dumps(extracted, indent=2)
         except json.JSONDecodeError:
             return "❌ JSON Parsing Fehler", json_text