Update app.py

app.py

@@ -4,18 +4,18 @@ import torch.nn as nn
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 def main():
-    # Get Hugging Face token from environment variable
+    # Get Hugging Face token from env
     auth_token = os.environ.get("HF_TOKEN")
-    if auth_token is None:
-        raise ValueError("Please set your Hugging Face token in the environment variable HF_TOKEN")
-
-    # Model ID
-    model_id = "google/gemma-3-1b-it"
+    if not auth_token:
+        raise ValueError("Please set HF_TOKEN environment variable")
 
     # Device
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    # Wrapper class
+    # Model ID
+    model_id = "google/gemma-3-1b-it"
+
+    # Wrapper
     class GemmaWrapper(nn.Module):
         def __init__(self, model_id, token):
             super().__init__()
@@ -28,7 +28,7 @@ def main():
         def forward(self, input_ids, attention_mask):
             return self.model(input_ids=input_ids, attention_mask=attention_mask, use_cache=False).logits
 
-    # Load model and tokenizer
+    # Load model & tokenizer
     model = GemmaWrapper(model_id, auth_token)
     tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=auth_token)
     tokenizer.pad_token = tokenizer.eos_token
@@ -41,10 +41,10 @@ def main():
 
     # Forward pass
     with torch.no_grad():
-        logits = model(input_ids=input_ids, attention_mask=attention_mask)
+        logits = model(input_ids, attention_mask)
 
     print("Logits shape:", logits.shape)
-    print("Sample logits:", logits[0, :5, :5])  # show small slice
+    print("Sample logits:", logits[0, :5, :5])
 
 if __name__ == "__main__":
     main()