Add rl example

client_tools/client_generator.py

@@ -240,7 +240,6 @@ def build_imports() -> Tuple[List[str], str]:
             if typing_imports:
                 lines.append(f"from typing import {', '.join(sorted(typing_imports))}")
             lines.extend([
-                "from twinkle_client.http import TWINKLE_SERVER_URL",
                 "from twinkle_client.http import http_post, heartbeat_manager",
             ])
             lines.extend(sorted(twinkle_imports))
@@ -447,7 +446,6 @@ def generate_models():
 
     model_code = AUTO_GEN_WARNING + '''from typing import Any, Optional, Union, Type, Dict, Literal, List
 import uuid
-from twinkle_client.http import TWINKLE_SERVER_URL
 from twinkle_client.http import http_post, heartbeat_manager
 from twinkle import DeviceMesh
 from twinkle.data_format import InputFeature, Trajectory
@@ -724,15 +722,10 @@ def generate_samplers():
     client_module_path.mkdir(parents=True, exist_ok=True)
 
     sampler_code = AUTO_GEN_WARNING + '''from typing import Any, Optional, List, Dict, Union
-import uuid
-from twinkle_client.http import TWINKLE_SERVER_URL
 from twinkle_client.http import http_post, heartbeat_manager
 from twinkle.sampler.base import Sampler
-from twinkle.sampler.types import SamplingParams, SampleResponse
-from twinkle import DeviceMesh
 from peft import PeftConfig
 from twinkle.data_format import Trajectory, InputFeature
-import json
 
 
 class vLLMSampler(Sampler):
@@ -750,13 +743,12 @@ def __init__(self, model_id: str, **kwargs):
         self.adapter_name = None
         if '://' in model_id:
             model_id = model_id.split('://')[1]
-        self.server_url = f'{self.server_url}/models/{model_id}'
+        self.server_url = f'{self.server_url}/samplers/{model_id}'
         response = http_post(
             url=f'{self.server_url}/create',
             json_data=kwargs
         )
         response.raise_for_status()
-        return response.json()
 
     def _send_adapter_heartbeat(self):
         """Internal method to send adapter heartbeat."""
@@ -799,25 +791,34 @@ def sample(
         self,
         inputs: Union[List[Trajectory], List[InputFeature]],
         sampling_params: Optional[Dict[str, Any]] = None,
-        adapter_name: str = ''
-    ) -> SampleResponse:
+        adapter_name: str = '',
+        adapter_uri: Optional[str] = None,
+        num_samples: int = 1,
+    ) -> Dict[str, Any]:
         """Sample from the model.
 
         Args:
             inputs: List of Trajectory or InputFeature to sample from.
             sampling_params: Sampling parameters dict.
-            adapter_name: Adapter name.
+            adapter_name: Adapter name for LoRA inference.
+            adapter_uri: Adapter URI (twinkle:// path or local path) for LoRA inference.
+            num_samples: Number of completions to generate per prompt.
 
         Returns:
-            SampleResponse with sampled sequences.
+            Dict with 'sequences' list, each containing tokens, logprobs, stop_reason.
         """
+        json_data = {
+            'inputs': inputs,
+            'sampling_params': sampling_params,
+            'adapter_name': adapter_name,
+            'num_samples': num_samples,
+        }
+        if adapter_uri is not None:
+            json_data['adapter_uri'] = adapter_uri
+
         response = http_post(
             url=f'{self.server_url}/sample',
-            json_data={
-                'inputs': inputs,
-                'sampling_params': sampling_params,
-                'adapter_name': adapter_name
-            }
+            json_data=json_data
         )
         response.raise_for_status()
         return response.json()