Oracle weight assigning update

Oracle/SmolLM.py

@@ -1,29 +1,38 @@
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 class SmolLM:
     def __init__(self, model_path="HuggingFaceTB/SmolLM2-1.7B-Instruct"):
         self.available = True
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
         try:
             print(f"[INFO] Loading Oracle tokenizer from {model_path}")
             self.tokenizer = AutoTokenizer.from_pretrained(model_path)
-            print(f"[INFO] Loading Oracle from {model_path}")
-            self.model = AutoModelForCausalLM.from_pretrained(model_path)
+            print(f"[INFO] Loading Oracle from {model_path} on {self.device}")
+            self.model = AutoModelForCausalLM.from_pretrained(model_path).to(self.device)
             print("[INFO] Oracle loaded successfully")
         except Exception as e:
             print(f"[ERROR] Failed to load model '{model_path}': {e}")
             self.available = False
 
-    def predict(self, prompt,max_length=512,max_new_tokens=150):
+    def predict(self, prompt, max_length=512, max_new_tokens=150):
         if not self.available:
             print("[WARN] Oracle unavailable, returning default weight 0.5")
             return "0.5"
         try:
-            print(f"[INFO] Generating response for prompt: {prompt[:100]}...", flush=True)
-            inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=max_length)
-            outputs = self.model.generate(**inputs, max_length=inputs["input_ids"].shape[1]+max_new_tokens,num_return_sequences=1)
+            # Use chat template as per documentation
+            messages = [{"role": "user", "content": prompt}]
+            inputs = self.tokenizer.apply_chat_template(messages, return_tensors="pt").to(self.device)
+            outputs = self.model.generate(
+                inputs,
+                max_new_tokens=max_new_tokens,
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True
+            )
             response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
             print(f"[INFO] Generated response: {response[:100]}...", flush=True)
             return response
         except Exception as e:
             print(f"[ERROR] Oracle has failed: {e}")
-            return "0.5"
+            return "0.5"