Update app.py

app.py

@@ -115,37 +115,36 @@ def load_model():
         print(f"✅ Loaded {MODEL_NAME} on {DEVICE}")
         
         # Test the model
-            test_text = "Test in Berlin."
-            test_template = '{"test_location": ""}'
-            test_template_formatted = json.dumps(json.loads(test_template), indent=4)
-            prompt = f"<|input|>\n### Template:\n{test_template_formatted}\n### Text:\n{test_text}\n\n<|output|>"
-            
-            # Create inputs with proper padding and truncation
-            inputs = tokenizer([prompt], return_tensors="pt", truncation=True, max_length=MAX_INPUT_LENGTH).to(DEVICE)
-            
-            # Generate output
-            with torch.no_grad():
-                outputs = model.generate(
-                    **inputs,
-                    max_new_tokens=50,
-                    temperature=0.0,
-                    do_sample=False
-                )
-                
-            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            
-            # Check if output contains expected text
-            if "<|output|>" in result and "Berlin" in result:
-                return "✅ Modell erfolgreich geladen und getestet!"
+        test_text = "Test in Berlin."
+        test_template = '{"test_location": ""}'
+        test_template_formatted = json.dumps(json.loads(test_template), indent=4)
+        prompt = f"<|input|>\n### Template:\n{test_template_formatted}\n### Text:\n{test_text}\n\n<|output|>"
+        
+        # Create inputs with proper padding and truncation
+        inputs = tokenizer([prompt], return_tensors="pt", truncation=True, max_length=MAX_INPUT_LENGTH).to(DEVICE)
+        
+        # Generate output
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=50,
+                temperature=0.0,
+                do_sample=False
+            )
             
+        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Check if output contains expected text
+        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:
         import traceback
         trace = traceback.format_exc()
         print(f"Error loading model: {e}\n{trace}")
         return f"❌ Fehler beim Laden des Modells: {str(e)}"
-        
 @spaces.GPU
 def extract_info(template, text):
     global tokenizer, model