Spaces:
No application file
No application file
File size: 4,637 Bytes
9529bc2 | 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 | from abc import ABC, abstractmethod
from typing import Optional
import json
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from ..core.robot_design import RobotDesign, RobotDesignBuilder
class LLMStrategy(ABC):
@abstractmethod
def generate_design(self, prompt: str) -> RobotDesign:
pass
class LocalLLMStrategy(LLMStrategy):
def __init__(self, model_id: str = "distilgpt2"):
self.model_id = model_id
self.tokenizer = None
self.model = None
self._load_model()
def _load_model(self):
try:
self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, local_files_only=True)
self.model = AutoModelForCausalLM.from_pretrained(self.model_id, local_files_only=True)
except Exception as e:
print(f"Failed to load model locally: {e}")
print("Downloading model...")
self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
self.model = AutoModelForCausalLM.from_pretrained(self.model_id)
def generate_design(self, prompt: str) -> RobotDesign:
formatted_prompt = f"""You are a robot design expert. Create a robot design based on the following requirements:
{prompt}
Return ONLY a JSON object with the following structure, no other text:
{{
"body": {{
"mass": 10.0,
"dimensions": [1.0, 1.0, 0.5]
}},
"wheels": [
{{
"radius": 0.2,
"width": 0.1,
"position": [0.5, 0.5, 0.2],
"friction": 0.8
}},
{{
"radius": 0.2,
"width": 0.1,
"position": [-0.5, 0.5, 0.2],
"friction": 0.8
}},
{{
"radius": 0.2,
"width": 0.1,
"position": [0.5, -0.5, 0.2],
"friction": 0.8
}},
{{
"radius": 0.2,
"width": 0.1,
"position": [-0.5, -0.5, 0.2],
"friction": 0.8
}}
],
"motors": [
{{
"max_force": 100.0,
"max_velocity": 10.0
}},
{{
"max_force": 100.0,
"max_velocity": 10.0
}},
{{
"max_force": 100.0,
"max_velocity": 10.0
}},
{{
"max_force": 100.0,
"max_velocity": 10.0
}}
]
}}"""
inputs = self.tokenizer(formatted_prompt, return_tensors="pt", truncation=True)
outputs = self.model.generate(
**inputs,
max_new_tokens=500,
num_return_sequences=1,
temperature=0.3,
do_sample=True,
top_p=0.9,
pad_token_id=self.tokenizer.eos_token_id
)
response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
try:
json_start = response.find('{')
json_end = response.rfind('}') + 1
if json_start == -1 or json_end == 0:
return RobotDesignBuilder.get_default_design()
json_str = response[json_start:json_end]
design_dict = json.loads(json_str)
# Validate required fields
required_fields = ["body", "wheels", "motors"]
for field in required_fields:
if field not in design_dict:
return RobotDesignBuilder.get_default_design()
# Create design using builder
builder = RobotDesignBuilder()
builder.set_body(
design_dict["body"]["mass"],
design_dict["body"]["dimensions"]
)
for wheel in design_dict["wheels"]:
builder.add_wheel(
wheel["radius"],
wheel["width"],
wheel["position"],
wheel["friction"]
)
for motor in design_dict["motors"]:
builder.add_motor(
motor["max_force"],
motor["max_velocity"]
)
return builder.build()
except Exception as e:
print(f"Failed to parse LLM response: {e}")
return RobotDesignBuilder.get_default_design()
class LLMFactory:
@staticmethod
def create_llm(llm_type: str = "local", model_id: Optional[str] = None) -> LLMStrategy:
if llm_type.lower() == "local":
return LocalLLMStrategy(model_id or "distilgpt2")
else:
raise ValueError(f"Unknown LLM type: {llm_type}") |