Update app.py

app.py

@@ -9,46 +9,51 @@ import uuid
 tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
 model = GPT2LMHeadModel.from_pretrained('gpt2')
 from datasets import load_dataset
-dataset = load_dataset("visionlab/block-towers-10k-3s-trajectory-scale1", split='train')
 class SecondLifeNavigator:
     def __init__(self):
-        # Load the dataset, focusing on the 'data' column if that's where the positions are stored.
-        self.dataset = self.load_dataset()
-        
-    def load_dataset(self):
-        # Load a specific split of the dataset, for example, 'train'.
-        
+        self.dataset = self.load_dataset("visionlab/block-towers-10k-3s-trajectory-scale1", split='train')
+        self.position_lookup = self.preprocess_data()
+
+    def load_dataset(self, dataset_name, split='train'):
+        dataset = load_dataset(dataset_name, split=split)
         return dataset
+    
+    def preprocess_data(self):
+        # Initialize an empty dictionary for the lookup table
+        position_lookup = {}
 
-    def determine_action_sequence(self, current_state):
-        # Assuming current_state can be used to derive current XYZ coordinates
-        current_xyz = [0, 0, 0]  # Example values, replace with actual extraction logic
-        
-        nearest_position = self.find_nearest_position(current_xyz)
-        if nearest_position:
-            # Construct action sequence based on nearest position
-            action_sequence = f"Move to position {nearest_position['x']}, {nearest_position['y']}, {nearest_position['z']}"
-        else:
-            action_sequence = "No suitable position found."
-        return action_sequence
-        
-    def find_nearest_position(self, current_xyz):
+        # Example: Creating a grid of positions for simplification
+        # This grid should be adjusted based on the expected range of xyz values
+        for x in np.arange(-100, 100, 1):  # Adjust the range and step as needed
+            for y in np.arange(-100, 100, 1):
+                for z in np.arange(-100, 100, 1):
+                    nearest_position = self.find_nearest_position([x, y, z], preprocess=True)
+                    position_lookup[(x, y, z)] = nearest_position
+
+        return position_lookup
+    
+    def find_nearest_position(self, current_xyz, preprocess=False):
         min_distance = float('inf')
         nearest_position = None
-    
-        # Iterate over items in the dataset. Assuming each item is a dictionary that follows the structure you provided.
+        
+        # Assuming 'trajectory' data contains the relevant positions during preprocessing
+        # If not preprocessing, use the lookup
+        if not preprocess:
+            # Here you access the pre-calculated nearest position from the lookup
+            # Adjust rounding or grid matching based on how the preprocess_data was implemented
+            return self.position_lookup.get(tuple(map(int, current_xyz)), None)
+        
         for item in self.dataset:
-            # Directly accessing the 'trajectory' data assuming 'item' is structured as described
-            trajectory_data = item['data']['trajectory']  # Adjust this if 'data' refers to something else
-    
+            trajectory_data = item['data']['trajectory']  # Adjust based on the actual structure
+            
             for record in trajectory_data:
-                for position in record['data']:  # Now we are in the 'data' list inside 'trajectory'
+                for position in record['data']:
                     xyz = position['xyz']
                     distance = np.sqrt((xyz[0] - current_xyz[0])**2 + (xyz[1] - current_xyz[1])**2 + (xyz[2] - current_xyz[2])**2)
                     if distance < min_distance:
                         min_distance = distance
                         nearest_position = xyz
-    
+
         return nearest_position
         
     def send_command_to_corrade(self, corrade_endpoint, command, parameters):