File size: 5,469 Bytes
9abace2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | """
PyTorch Tools
Part of SOVEREIGN PYTHON LLM ENGINE
Enables agents to use PyTorch for tensor operations and model inference.
"""
from typing import Any
import json
async def tensor_operation(params: dict[str, Any]) -> dict[str, Any]:
"""
Execute PyTorch tensor operation.
Args:
params: {
"operation": str, # "add", "matmul", "softmax", etc
"tensor_a": list, # Input tensor data
"tensor_b": list | None, # Second tensor (optional)
"dim": int | None # Dimension for operations
}
Returns:
dict with "result" (tensor as list) and "shape"
"""
try:
import torch
except ImportError:
return {"error": "PyTorch not installed. Run: pip install torch"}
operation = params.get("operation")
tensor_a_data = params.get("tensor_a")
tensor_b_data = params.get("tensor_b")
dim = params.get("dim")
if not operation or not tensor_a_data:
return {"error": "operation and tensor_a required"}
try:
# Convert to tensors
tensor_a = torch.tensor(tensor_a_data)
tensor_b = torch.tensor(tensor_b_data) if tensor_b_data else None
# Execute operation
if operation == "add":
result = tensor_a + tensor_b if tensor_b is not None else tensor_a
elif operation == "matmul":
if tensor_b is None:
return {"error": "tensor_b required for matmul"}
result = torch.matmul(tensor_a, tensor_b)
elif operation == "softmax":
result = torch.softmax(tensor_a, dim=dim if dim is not None else -1)
elif operation == "relu":
result = torch.relu(tensor_a)
elif operation == "mean":
result = torch.mean(tensor_a, dim=dim)
elif operation == "sum":
result = torch.sum(tensor_a, dim=dim)
elif operation == "transpose":
result = torch.transpose(tensor_a, dim0=0, dim1=1)
else:
return {"error": f"Unknown operation: {operation}"}
return {
"result": result.tolist(),
"shape": list(result.shape),
"dtype": str(result.dtype)
}
except Exception as e:
return {"error": f"PyTorch operation failed: {str(e)}"}
async def load_model(params: dict[str, Any]) -> dict[str, Any]:
"""
Load PyTorch model from file.
Args:
params: {
"model_path": str, # Path to model file (.pt, .pth)
"device": str | None # "cpu" or "cuda"
}
Returns:
dict with "model_id" and "device"
"""
try:
import torch
except ImportError:
return {"error": "PyTorch not installed"}
model_path = params.get("model_path")
device = params.get("device", "cpu")
if not model_path:
return {"error": "model_path required"}
try:
# Check if CUDA available
if device == "cuda" and not torch.cuda.is_available():
device = "cpu"
print("CUDA not available, using CPU")
# Load model
model = torch.load(model_path, map_location=device)
# Store model in global cache (simple implementation)
# In production, use proper model registry
model_id = f"model_{id(model)}"
return {
"model_id": model_id,
"device": device,
"success": True
}
except Exception as e:
return {"error": f"Failed to load model: {str(e)}"}
async def model_inference(params: dict[str, Any]) -> dict[str, Any]:
"""
Run inference on loaded PyTorch model.
Args:
params: {
"model_id": str, # Model ID from load_model
"input": list, # Input tensor data
"batch_size": int | None
}
Returns:
dict with "output" (predictions) and "shape"
"""
try:
import torch
except ImportError:
return {"error": "PyTorch not installed"}
# NOTE: This is a simplified implementation
# In production, would use proper model registry
return {
"error": "Model inference requires model registry (TODO)",
"note": "Use tensor_operation for basic PyTorch ops"
}
async def check_cuda(params: dict[str, Any]) -> dict[str, Any]:
"""
Check CUDA availability and device info.
Returns:
dict with CUDA status and device info
"""
try:
import torch
except ImportError:
return {"error": "PyTorch not installed"}
cuda_available = torch.cuda.is_available()
result = {
"cuda_available": cuda_available,
"cuda_device_count": torch.cuda.device_count() if cuda_available else 0,
"pytorch_version": torch.__version__
}
if cuda_available:
result["cuda_version"] = torch.version.cuda
result["cudnn_version"] = torch.backends.cudnn.version()
result["devices"] = [
{
"index": i,
"name": torch.cuda.get_device_name(i),
"memory_total": torch.cuda.get_device_properties(i).total_memory,
"capability": torch.cuda.get_device_capability(i)
}
for i in range(torch.cuda.device_count())
]
return result
|