Remove logging and disable field model

app.py

@@ -3,7 +3,7 @@ from gradio_client import Client
 from gradio_client.exceptions import AppError
 import frontmatter
 import os
-#import spaces
+import spaces
 import torch
 import logging
 from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -16,18 +16,6 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# Enable transformers logging
-transformers_logging.set_verbosity_debug()
-transformers_logging.enable_default_handler()
-transformers_logging.enable_explicit_format()
-
-# Enable accelerate and torch logging
-logging.getLogger("accelerate").setLevel(logging.DEBUG)
-logging.getLogger("torch").setLevel(logging.DEBUG)
-logging.getLogger("spaces").setLevel(logging.DEBUG)
-logging.getLogger("spaces.zero").setLevel(logging.DEBUG)
-logging.getLogger("transformers").setLevel(logging.DEBUG)
-
 import huggingface_hub
 
 import prep_decompiled
@@ -43,26 +31,18 @@ tokenizer = AutoTokenizer.from_pretrained("bigcode/starcoderbase-3b")
 vardecoder_model = AutoModelForCausalLM.from_pretrained(
     "ejschwartz/resym-vardecoder", 
     torch_dtype=torch.bfloat16,
+    device_map="auto",
 )
 print("Loaded vardecoder model successfully.")
 
-print(f"Model device: {next(vardecoder_model.parameters()).device}")
-print(f"Model dtype: {next(vardecoder_model.parameters()).dtype}")
-print(f"Model is meta: {next(vardecoder_model.parameters()).is_meta}")
-print(f"Model parameters: {sum(p.numel() for p in vardecoder_model.parameters() if p.requires_grad):,}")
-
-# Check memory after first model
-print(f"GPU memory after vardecoder:")
-print(f"Allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
-print(f"Reserved: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")
-
 logger.info("Loading fielddecoder model...")
 
-fielddecoder_model = AutoModelForCausalLM.from_pretrained(
-    "ejschwartz/resym-fielddecoder", 
-    torch_dtype=torch.bfloat16,
-)
-logger.info("Successfully loaded fielddecoder model")
+fielddecoder_model = None
+#fielddecoder_model = AutoModelForCausalLM.from_pretrained(
+#    "ejschwartz/resym-fielddecoder", 
+#    torch_dtype=torch.bfloat16,
+#)
+#logger.info("Successfully loaded fielddecoder model")
 
 make_gradio_client = lambda: Client("https://ejschwartz-resym-field-helper.hf.space/")
 
@@ -155,23 +135,26 @@ def infer(code):
             :, : MAX_CONTEXT_LENGTH - MAX_NEW_TOKENS
         ]
 
-        field_output = fielddecoder_model.generate(
-            input_ids=field_input_ids,
-            max_new_tokens=MAX_NEW_TOKENS,
-            num_beams=4,
-            num_return_sequences=1,
-            do_sample=False,
-            early_stopping=False,
-            pad_token_id=0,
-            eos_token_id=0,
-        )[0]
-        field_output = tokenizer.decode(
-            field_output[field_input_ids.size(1) :],
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=True,
-        )
-
-        field_output = fields[0] + ":" + field_output
+        if fielddecoder_model is None:
+            field_output = "TEMPORARILY DISABLED"
+        else:
+            field_output = fielddecoder_model.generate(
+                input_ids=field_input_ids,
+                max_new_tokens=MAX_NEW_TOKENS,
+                num_beams=4,
+                num_return_sequences=1,
+                do_sample=False,
+                early_stopping=False,
+                pad_token_id=0,
+                eos_token_id=0,
+            )[0]
+            field_output = tokenizer.decode(
+                field_output[field_input_ids.size(1) :],
+                skip_special_tokens=True,
+                clean_up_tokenization_spaces=True,
+            )
+
+            field_output = fields[0] + ":" + field_output
     var_output = first_var + ":" + var_output
     fieldstring = ", ".join(fields)
     return var_output, field_output, varstring, fieldstring