Update main.py

main.py

@@ -65,26 +65,27 @@ def predict_single(code):
         preds = torch.sigmoid(outputs.logits).cpu().numpy()
         cpu_time, memory_usage = preds[0]
 
+        # Invert values so higher = better
         return {
-            "cpu_time": round(float(cpu_time), 4),
-            "memory_usage": round(float(memory_usage), 4)
+            "cpu_efficiency": round(1.0 - float(cpu_time), 4),
+            "memory_efficiency": round(1.0 - float(memory_usage), 4)
         }
 
     except Exception as e:
         print(f"Single prediction error: {e}")
-        return {"cpu_time": 0.0, "memory_usage": 0.0}
+        return {"cpu_efficiency": 0.0, "memory_efficiency": 0.0}
 
 def predict_with_chunking(code, chunk_size=400, overlap=50):
     try:
         if not code or len(code.strip()) == 0:
-            return {"cpu_time": 0.0, "memory_usage": 0.0}
+            return {"cpu_efficiency": 0.0, "memory_efficiency": 0.0}
 
         tokens = tokenizer.encode(code, add_special_tokens=False)
         if len(tokens) <= 450:
             return predict_single(code)
 
-        max_cpu_time = 0.0
-        max_memory_usage = 0.0
+        max_cpu_efficiency = 0.0
+        max_memory_efficiency = 0.0
 
         for start in range(0, len(tokens), chunk_size - overlap):
             end = min(start + chunk_size, len(tokens))
@@ -93,20 +94,20 @@ def predict_with_chunking(code, chunk_size=400, overlap=50):
 
             if chunk_code.strip():
                 result = predict_single(chunk_code)
-                max_cpu_time = max(max_cpu_time, result["cpu_time"])
-                max_memory_usage = max(max_memory_usage, result["memory_usage"])
+                max_cpu_efficiency = max(max_cpu_efficiency, result["cpu_efficiency"])
+                max_memory_efficiency = max(max_memory_efficiency, result["memory_efficiency"])
 
             if end >= len(tokens):
                 break
 
         return {
-            "cpu_time": round(max_cpu_time, 4),
-            "memory_usage": round(max_memory_usage, 4)
+            "cpu_efficiency": round(max_cpu_efficiency, 4),
+            "memory_efficiency": round(max_memory_efficiency, 4)
         }
 
     except Exception as e:
         print(f"Chunking prediction error: {e}")
-        return {"cpu_time": 0.0, "memory_usage": 0.0}
+        return {"cpu_efficiency": 0.0, "memory_efficiency": 0.0}
 
 @app.route("/", methods=['GET'])
 def home():