Initial commit

inference.py

@@ -0,0 +1,40 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import lightgbm as lgb
+import numpy as np
+
+class LinearProbeWrapper:
+    def __init__(self, model_name):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(model_name, output_hidden_states=True)
+        self.linear_probe = lgb.Booster(model_file='lgb_layer_1.txt')
+
+
+    def __call__(self, inputs):
+        prompt = inputs.get("inputs", "")
+        inputs_tensors = self.tokenizer(prompt, return_tensors="pt")
+
+        with torch.no_grad():
+            # Get hidden states
+            outputs = self.model(**inputs_tensors)
+            # Access the hidden states (residual stream) - adjust index as needed
+            layer_1_hidden_state = outputs.hidden_states[1]
+            layer_1_hidden_state_np = layer_1_hidden_state.cpu().numpy().copy()
+
+            #Predict given hidden state
+            y_pred = self.linear_probe.predict(layer_1_hidden_state_np[0]);
+            y_pred_class = np.argmax(y_pred)
+
+            # Generate text
+            generation_output = self.model.generate(
+                **inputs_tensors, max_length=50, num_return_sequences=1
+            )
+            generated_text = self.tokenizer.decode(generation_output[0], skip_special_tokens=True)
+
+        return {
+            "generated_text": generated_text,
+            "probe_output": y_pred[y_pred_class]
+        }
+
+def model_fn():
+    return LinearProbeWrapper("mistralai/Mistral-7B-Instruct-v0.2")  # Replace with desired model

requirements.txt

@@ -0,0 +1,4 @@
+torch
+transformers
+lightgbm
+numpy