Update app.py

app.py

@@ -1,78 +1,35 @@
-import os
 import torch
-import torch.nn as nn
-from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# --- ToyModel placeholder ---
-# Replace this with your actual ToyModel import if it exists
-class ToyModel:
-    def __init__(self, model, inputs, dynamic_shapes=None):
-        self.model = model
-        self.inputs = inputs
-        self.dynamic_shapes = dynamic_shapes
-
-    def predict(self, input_ids, attention_mask):
-        with torch.no_grad():
-            logits = self.model(input_ids=input_ids, attention_mask=attention_mask)
-        return logits
-
-    def evaluation(self):
-        print("Evaluation complete.")
-
-# ----------------------------
-
-def main():
-    # Hugging Face token from environment
-    auth_token = os.environ.get("HF_TOKEN")
-    if not auth_token:
-        raise ValueError("Please set HF_TOKEN environment variable")
-
-    # Device
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    # Model ID
-    model_id = "google/gemma-3-1b-it"
-
-    # Wrapper
-    class GemmaWrapper(nn.Module):
-        def __init__(self, model_id, token):
-            super().__init__()
-            self.model = AutoModelForCausalLM.from_pretrained(
-                model_id,
-                torch_dtype=torch.float32,
-                use_auth_token=token
-            ).to(device).eval()
-        
-        def forward(self, input_ids, attention_mask):
-            return self.model(input_ids=input_ids, attention_mask=attention_mask, use_cache=False).logits
-
-    # 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
-
-    # Example input
-    sentences = ["Hello"]
-    tokens = tokenizer(sentences, return_tensors="pt", padding=True, truncation=True)
-    input_ids = tokens["input_ids"].to(device)
-    attention_mask = tokens["attention_mask"].to(device)
-
-    # Dynamic shapes (optional, can be used by ToyModel)
-    batch_dim = 1
-    seq_dim = input_ids.shape[1]
-    dynamic_shapes = {
-        "input_ids": {0: batch_dim, 1: seq_dim},
-        "attention_mask": {0: batch_dim, 1: seq_dim},
-    }
-
-    # --- ToyModel usage ---
-    ir = ToyModel(model, (input_ids, attention_mask), dynamic_shapes=dynamic_shapes)
-    io_data = ir.predict(input_ids, attention_mask)
-    ir.evaluation()
-
-    # Output
-    print("Predicted logits shape:", io_data.shape)
-    print("Sample logits:", io_data[0, :5, :5])
-
-if __name__ == "__main__":
-    main()
+# Utility to log tensor info
+def log_tensor(name, x):
+    print(f"--- {name} ---")
+    print(f"shape: {x.shape}, dtype: {x.dtype}, device: {x.device}")
+    print(f"min: {x.min().item():.6f}, max: {x.max().item():.6f}, mean: {x.mean().item():.6f}, sum: {x.sum().item():.6f}")
+    print(f"full tensor:\n{x}\n")
+
+# Simple function
+def g(x, y):
+    log_tensor("g input x", x)
+    log_tensor("g input y", y)
+    z = x + y
+    log_tensor("g output z", z)
+    return z
+
+# Compiled function
+@torch.compile(backend="eager")
+def f(x):
+    log_tensor("f input x", x)
+    x = torch.sin(x)
+    log_tensor("f after torch.sin(x)", x)
+    x = g(x, x)
+    log_tensor("f after g(x, x)", x)
+    return x
+
+# Example input
+x = torch.ones(3, 3, dtype=torch.float32)
+log_tensor("original input x", x)
+
+# Run compiled function
+out = f(x)
+
+log_tensor("final output", out)