Create policy.py

policy.py

@@ -0,0 +1,18 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+def load_policy_model():
+    model_name = "microsoft/phi-2"
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        device_map="auto",
+        torch_dtype=torch.float32
+    )
+
+    # Attach simple optimizer
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
+    model.optimizer = optimizer
+
+    return model, tokenizer