Loading vagen/inference/run_inference.py +16 −63 Original line number Diff line number Diff line Loading @@ -104,58 +104,22 @@ def run_inference_batch( return all_results def log_results_to_wandb(results: List[Dict], global_step: int = 0) -> None: """Log results to wandb using the same format as training.""" # Log metrics by config_id (same as training) metrics = log_metrics_by_config_id(results, mode='val') # Use 'val' to match training prefix """Log results to wandb without any aggregation.""" # Log raw metrics for each environment metrics = log_metrics_by_config_id(results, mode='val') wandb.log(metrics) # Log generation table (same as training) # Log generation table (unchanged) generations_to_log = 10 # You can make this configurable maybe_log_val_generations_to_wandb(results, generations_to_log, global_step) # Log overall statistics success_count = sum(1 for r in results if r['metrics'].get('success', 0) > 0) done_count = sum(1 for r in results if r['metrics'].get('done', 0) > 0) total_score = sum(r['metrics'].get('score', 0) for r in results) avg_steps = sum(r['metrics'].get('step', 0) for r in results) / len(results) if results else 0 # If you still want some overall summary, just count things wandb.log({ "val/success_rate": success_count / len(results) if results else 0, "val/completion_rate": done_count / len(results) if results else 0, "val/average_score": total_score / len(results) if results else 0, "val/average_steps": avg_steps "val/total_environments": len(results), "val/num_successful": sum(1 for r in results if r['metrics'].get('success', 0) > 0), "val/num_done": sum(1 for r in results if r['metrics'].get('done', 0) > 0), }) def save_results(results: List[Dict], output_dir: str) -> None: """Save results to disk.""" os.makedirs(output_dir, exist_ok=True) # Custom JSON encoder to handle NumPy types class NumpyEncoder(json.JSONEncoder): def default(self, obj): import numpy as np if isinstance(obj, np.integer): return int(obj) elif isinstance(obj, np.floating): return float(obj) elif isinstance(obj, np.bool_): return bool(obj) elif isinstance(obj, np.ndarray): return obj.tolist() return super(NumpyEncoder, self).default(obj) # Save raw results results_file = os.path.join(output_dir, "results.json") with open(results_file, "w") as f: json.dump(results, f, indent=2, cls=NumpyEncoder) # Save summary summary = calculate_aggregate_metrics(results) summary_file = os.path.join(output_dir, "summary.json") with open(summary_file, "w") as f: json.dump(summary, f, indent=2, cls=NumpyEncoder) logger.info(f"Results saved to {output_dir}") def main(): """Main entry point for inference.""" Loading Loading @@ -209,35 +173,24 @@ def main(): max_steps=inference_config.get('max_steps', 10) ) # Log results to wandb (aligned with training format) # Log results to wandb if inference_config.get('use_wandb', True): log_results_to_wandb(results, global_step=0) # Save results output_dir = os.path.join( inference_config.get('output_dir', 'inference_outputs'), model_name, datetime.now().strftime("%Y%m%d_%H%M%S") ) save_results(results, output_dir) # Print summary # Print summary (no saving to disk) print(f"\n===== Results for {model_name} =====") print(f"Total environments: {len(results)}") success_count = sum(1 for r in results if r['metrics'].get('success', 0) > 0) done_count = sum(1 for r in results if r['metrics'].get('done', 0) > 0) total_score = sum(r['metrics'].get('score', 0) for r in results) print(f"Success rate: {success_count / len(results) * 100:.1f}%") print(f"Completion rate: {done_count / len(results) * 100:.1f}%") print(f"Average score: {total_score / len(results):.4f}") print(f"Successful: {success_count}") print(f"Completed: {done_count}") # Print by config_id results (like training) metrics_by_config = log_metrics_by_config_id(results, mode='val') print("\nMetrics by config:") for metric_name, value in metrics_by_config.items(): print(f" {metric_name}: {value:.4f}") # Print individual results print("\nIndividual results:") for i, result in enumerate(results): print(f"Environment {i}: score={result['metrics'].get('score', 0)}, done={result['metrics'].get('done', 0)}, steps={result['metrics'].get('step', 0)}") except Exception as e: logger.error(f"Error during inference for model {model_name}: {str(e)}") Loading vagen/inference/utils/logging.py +52 −75 Original line number Diff line number Diff line Loading @@ -20,10 +20,7 @@ class ValidationTableManager: generations_to_log: int, global_steps: int ) -> None: """ Log a table of validation samples with multiple images per sample to wandb. Matches the training code's _maybe_log_val_generations_to_wandb function. """ """Log a table of validation samples with each example in its own row.""" if generations_to_log == 0: return Loading @@ -31,7 +28,7 @@ class ValidationTableManager: logger.warning('`val_generations_to_log_to_wandb` is set, but wandb is not initialized') return # Extract data from results exactly like training code # Extract data from results inputs = [] outputs = [] scores = [] Loading @@ -43,79 +40,67 @@ class ValidationTableManager: scores.append(item['metrics']['score']) images.append(item['image_data']) # Handle the case where images might not be provided if images is None or all(img is None for img in images): samples = list(zip(inputs, outputs, scores)) has_images = False max_images_per_sample = 0 else: samples = list(zip(inputs, outputs, scores, images)) has_images = True # Check if we have images has_images = any(img_list for img_list in images) # Find maximum number of images in any sample if has_images: max_images_per_sample = max( len(img_list) if isinstance(img_list, (list, tuple)) else 1 len(img_list) if img_list else 0 for img_list in images ) else: max_images_per_sample = 0 # Sort and shuffle exactly like training samples.sort(key=lambda x: x[0]) # Sort by input text # Create samples if has_images: samples = list(zip(inputs, outputs, scores, images)) else: samples = list(zip(inputs, outputs, scores)) # Use fixed random seed for deterministic shuffling rng = np.random.RandomState() # No seed specified in training code # Sort and shuffle for consistency samples.sort(key=lambda x: x[0]) # Sort by input text rng = np.random.RandomState() rng.shuffle(samples) # Take first N samples after shuffling # Take first N samples samples = samples[:generations_to_log] # Create column names for all samples # Create columns for the table if has_images: columns = ["step"] for i in range(len(samples)): columns.extend([f"input_{i+1}", f"output_{i+1}", f"score_{i+1}"]) columns.extend([f"image_{i+1}_{j+1}" for j in range(max_images_per_sample)]) columns = ["step", "input", "output", "score"] + [f"image_{i+1}" for i in range(max_images_per_sample)] else: columns = ["step"] + sum([[f"input_{i+1}", f"output_{i+1}", f"score_{i+1}"] for i in range(len(samples))], []) if self.validation_table is None: # Initialize the table on first call self.validation_table = wandb.Table(columns=columns) # Create a new table with same columns and existing data new_table = wandb.Table(columns=columns, data=self.validation_table.data) columns = ["step", "input", "output", "score"] # Add new row with all data row_data = [] row_data.append(global_steps) # Create table table = wandb.Table(columns=columns) # Add each sample as a separate row for sample in samples: if has_images: input_text, output_text, score, sample_images = sample row_data.extend([input_text, output_text, score]) # Handle if sample_images is a single image or list of images if not isinstance(sample_images, (list, tuple)): sample_images = [sample_images] # Convert each image to wandb.Image # Convert images to wandb.Image wandb_images = [] if sample_images: for img in sample_images: if img is not None: if not isinstance(img, wandb.Image): img = wandb.Image(img) wandb_images.append(img) # Pad with None if there are fewer images than max_images_per_sample wandb_images.extend([None] * (max_images_per_sample - len(wandb_images))) row_data.extend(wandb_images) # Pad with None if fewer images than max while len(wandb_images) < max_images_per_sample: wandb_images.append(None) # Add row table.add_data(global_steps, input_text, output_text, score, *wandb_images) else: input_text, output_text, score = sample row_data.extend([input_text, output_text, score]) new_table.add_data(*row_data) table.add_data(global_steps, input_text, output_text, score) # Update reference and log wandb.log({"val/generations": new_table}) self.validation_table = new_table # Log the table wandb.log({"val/generations": table}) # Global instance for maintaining table state across calls Loading @@ -141,29 +126,21 @@ def maybe_log_val_generations_to_wandb( def log_metrics_by_config_id( results: List[Dict[str, Any]], mode: str = 'inference' ) -> Dict[str, float]: ) -> Dict[str, Any]: """ Convert results to metrics dictionary with proper prefixes. Matches training code's log_rst_to_metrics_dict function. Convert results to metrics dictionary without any aggregation. Just logs raw metrics as they are. """ from collections import defaultdict metric_dict = {} # Group metrics by config_id metrics_by_config_id = defaultdict(dict) # dict of dict of list for item in results: # Simply iterate through results and create metric entries for i, item in enumerate(results): config_id = item["config_id"] env_id = item.get("env_id", f"env_{i}") # Log each metric as is, with a unique key including env_id for k, v in item["metrics"].items(): if k not in metrics_by_config_id[config_id]: metrics_by_config_id[config_id][k] = [] metrics_by_config_id[config_id][k].append(v) # Aggregate metrics for config_id, metrics in metrics_by_config_id.items(): for k, v in metrics.items(): metric_key = f'{mode}/{k}/{config_id}' metric_dict[metric_key] = np.mean(v) metric_key = f'{mode}/{config_id}/{env_id}/{k}' metric_dict[metric_key] = v return metric_dict No newline at end of file vagen/mllm_agent/model_interface/__init__.py +5 −0 Original line number Diff line number Diff line from .vllm import VLLMModelInterface, VLLMModelConfig from .openai import OpenAIModelInterface, OpenAIModelConfig REGISTERED_MODEL = { "vllm": { "model_cls": VLLMModelInterface, "config_cls": VLLMModelConfig, }, "openai": { "model_cls": OpenAIModelInterface, "config_cls": OpenAIModelConfig } } No newline at end of file vagen/mllm_agent/model_interface/openai/__init__.py 0 → 100644 +2 −0 Original line number Diff line number Diff line from .model import OpenAIModelInterface from .model_config import OpenAIModelConfig No newline at end of file vagen/mllm_agent/model_interface/openai/model.py 0 → 100644 +221 −0 Original line number Diff line number Diff line # vagen/mllm_agent/model_interface/openai/model.py import base64 import logging import re from typing import List, Dict, Any from concurrent.futures import ThreadPoolExecutor from openai import OpenAI from PIL import Image import io from vagen.mllm_agent.model_interface.base_model import BaseModelInterface from .model_config import OpenAIModelConfig logger = logging.getLogger(__name__) class OpenAIModelInterface(BaseModelInterface): """Model interface for OpenAI API with Qwen format compatibility.""" def __init__(self, config: OpenAIModelConfig): super().__init__(config) self.config = config # Initialize OpenAI client self.client = OpenAI( api_key=config.api_key, organization=config.organization, base_url=config.base_url ) # Thread pool for batch processing self.executor = ThreadPoolExecutor(max_workers=10) logger.info(f"Initialized OpenAI interface with model {config.model_name}") def generate(self, prompts: List[Any], **kwargs) -> List[Dict[str, Any]]: """Generate responses using OpenAI API.""" # Process prompts into OpenAI message format formatted_requests = [] for prompt in prompts: messages = self._convert_qwen_to_openai_format(prompt) formatted_requests.append(messages) # Make parallel API calls futures = [] for messages in formatted_requests: future = self.executor.submit( self._single_api_call, messages, **kwargs ) futures.append(future) # Collect results results = [] for future in futures: try: result = future.result() results.append(result) except Exception as e: logger.error(f"API call failed: {e}") results.append({ "text": f"Error: {str(e)}", "error": str(e) }) return results def _convert_qwen_to_openai_format(self, prompt: List[Dict]) -> List[Dict]: """ Convert Qwen format messages to OpenAI format. Qwen format: Text with <image> placeholders + separate multi_modal_data OpenAI format: Structured content array with text and image objects """ openai_messages = [] for message in prompt: role = message.get("role", "user") content = message.get("content", "") # Create OpenAI message structure openai_msg = { "role": role, "content": [] } # Handle multimodal content if "multi_modal_data" in message and "<image>" in content: # Extract images from multi_modal_data images = [] for key, values in message["multi_modal_data"].items(): if key == "<image>" or "image" in key.lower(): images.extend(values) # Split content by <image> placeholders parts = content.split("<image>") # Build content array alternating text and images for i, part in enumerate(parts): # Add text part if not empty if part.strip(): openai_msg["content"].append({ "type": "text", "text": part }) # Add image if available (except for last part) if i < len(parts) - 1 and i < len(images): image_data = self._process_image_for_openai(images[i]) openai_msg["content"].append({ "type": "image_url", "image_url": { "url": f"data:image/jpeg;base64,{image_data}" } }) else: # Text-only message openai_msg["content"].append({ "type": "text", "text": content }) openai_messages.append(openai_msg) return openai_messages def _process_image_for_openai(self, image: Any) -> str: """Convert image to base64 for OpenAI API.""" if isinstance(image, Image.Image): # Ensure RGB mode if image.mode != "RGB": image = image.convert("RGB") # Resize if too large to save tokens max_size = 1024 if max(image.size) > max_size: ratio = max_size / max(image.size) new_size = tuple(int(dim * ratio) for dim in image.size) image = image.resize(new_size, Image.Resampling.LANCZOS) buffered = io.BytesIO() image.save(buffered, format="JPEG", quality=85) return base64.b64encode(buffered.getvalue()).decode() elif isinstance(image, dict) and "__pil_image__" in image: from vagen.server.serial import deserialize_pil_image pil_image = deserialize_pil_image(image) return self._process_image_for_openai(pil_image) else: raise ValueError(f"Unsupported image type: {type(image)}") def _single_api_call(self, messages: List[Dict], **kwargs) -> Dict[str, Any]: """Make a single API call to OpenAI.""" try: response = self.client.chat.completions.create( model=self.config.model_name, messages=messages, max_tokens=kwargs.get("max_tokens", self.config.max_tokens), temperature=kwargs.get("temperature", self.config.temperature), presence_penalty=kwargs.get("presence_penalty", self.config.presence_penalty), frequency_penalty=kwargs.get("frequency_penalty", self.config.frequency_penalty), seed=kwargs.get("seed", self.config.seed) ) # Extract text response response_text = response.choices[0].message.content # Convert response back to Qwen format if needed # (OpenAI doesn't return images, so no conversion needed for output) return { "text": response_text, "usage": { "prompt_tokens": response.usage.prompt_tokens, "completion_tokens": response.usage.completion_tokens, "total_tokens": response.usage.total_tokens }, "finish_reason": response.choices[0].finish_reason } except Exception as e: logger.error(f"OpenAI API error: {e}") raise def format_prompt(self, messages: List[Dict[str, Any]]) -> str: """ Format prompt for compatibility. Since OpenAI uses structured messages, this returns a string representation of the messages for logging/debugging purposes. """ formatted = [] for msg in messages: role = msg.get("role", "user") content = msg.get("content", "") # Handle Qwen special tokens if present if role == "system": formatted.append(f"System: {content}") elif role == "user": formatted.append(f"User: {content}") elif role == "assistant": formatted.append(f"Assistant: {content}") return "\n".join(formatted) def get_model_info(self) -> Dict[str, Any]: """Get detailed information about the model.""" info = super().get_model_info() info.update({ "name": self.config.model_name, "type": "multimodal" if "vision" in self.config.model_name.lower() else "text", "supports_images": "vision" in self.config.model_name.lower() or "4o" in self.config.model_name, "max_tokens": self.config.max_tokens, "temperature": self.config.temperature, "config_id": self.config.config_id() }) return info No newline at end of file Loading
vagen/inference/run_inference.py +16 −63 Original line number Diff line number Diff line Loading @@ -104,58 +104,22 @@ def run_inference_batch( return all_results def log_results_to_wandb(results: List[Dict], global_step: int = 0) -> None: """Log results to wandb using the same format as training.""" # Log metrics by config_id (same as training) metrics = log_metrics_by_config_id(results, mode='val') # Use 'val' to match training prefix """Log results to wandb without any aggregation.""" # Log raw metrics for each environment metrics = log_metrics_by_config_id(results, mode='val') wandb.log(metrics) # Log generation table (same as training) # Log generation table (unchanged) generations_to_log = 10 # You can make this configurable maybe_log_val_generations_to_wandb(results, generations_to_log, global_step) # Log overall statistics success_count = sum(1 for r in results if r['metrics'].get('success', 0) > 0) done_count = sum(1 for r in results if r['metrics'].get('done', 0) > 0) total_score = sum(r['metrics'].get('score', 0) for r in results) avg_steps = sum(r['metrics'].get('step', 0) for r in results) / len(results) if results else 0 # If you still want some overall summary, just count things wandb.log({ "val/success_rate": success_count / len(results) if results else 0, "val/completion_rate": done_count / len(results) if results else 0, "val/average_score": total_score / len(results) if results else 0, "val/average_steps": avg_steps "val/total_environments": len(results), "val/num_successful": sum(1 for r in results if r['metrics'].get('success', 0) > 0), "val/num_done": sum(1 for r in results if r['metrics'].get('done', 0) > 0), }) def save_results(results: List[Dict], output_dir: str) -> None: """Save results to disk.""" os.makedirs(output_dir, exist_ok=True) # Custom JSON encoder to handle NumPy types class NumpyEncoder(json.JSONEncoder): def default(self, obj): import numpy as np if isinstance(obj, np.integer): return int(obj) elif isinstance(obj, np.floating): return float(obj) elif isinstance(obj, np.bool_): return bool(obj) elif isinstance(obj, np.ndarray): return obj.tolist() return super(NumpyEncoder, self).default(obj) # Save raw results results_file = os.path.join(output_dir, "results.json") with open(results_file, "w") as f: json.dump(results, f, indent=2, cls=NumpyEncoder) # Save summary summary = calculate_aggregate_metrics(results) summary_file = os.path.join(output_dir, "summary.json") with open(summary_file, "w") as f: json.dump(summary, f, indent=2, cls=NumpyEncoder) logger.info(f"Results saved to {output_dir}") def main(): """Main entry point for inference.""" Loading Loading @@ -209,35 +173,24 @@ def main(): max_steps=inference_config.get('max_steps', 10) ) # Log results to wandb (aligned with training format) # Log results to wandb if inference_config.get('use_wandb', True): log_results_to_wandb(results, global_step=0) # Save results output_dir = os.path.join( inference_config.get('output_dir', 'inference_outputs'), model_name, datetime.now().strftime("%Y%m%d_%H%M%S") ) save_results(results, output_dir) # Print summary # Print summary (no saving to disk) print(f"\n===== Results for {model_name} =====") print(f"Total environments: {len(results)}") success_count = sum(1 for r in results if r['metrics'].get('success', 0) > 0) done_count = sum(1 for r in results if r['metrics'].get('done', 0) > 0) total_score = sum(r['metrics'].get('score', 0) for r in results) print(f"Success rate: {success_count / len(results) * 100:.1f}%") print(f"Completion rate: {done_count / len(results) * 100:.1f}%") print(f"Average score: {total_score / len(results):.4f}") print(f"Successful: {success_count}") print(f"Completed: {done_count}") # Print by config_id results (like training) metrics_by_config = log_metrics_by_config_id(results, mode='val') print("\nMetrics by config:") for metric_name, value in metrics_by_config.items(): print(f" {metric_name}: {value:.4f}") # Print individual results print("\nIndividual results:") for i, result in enumerate(results): print(f"Environment {i}: score={result['metrics'].get('score', 0)}, done={result['metrics'].get('done', 0)}, steps={result['metrics'].get('step', 0)}") except Exception as e: logger.error(f"Error during inference for model {model_name}: {str(e)}") Loading
vagen/inference/utils/logging.py +52 −75 Original line number Diff line number Diff line Loading @@ -20,10 +20,7 @@ class ValidationTableManager: generations_to_log: int, global_steps: int ) -> None: """ Log a table of validation samples with multiple images per sample to wandb. Matches the training code's _maybe_log_val_generations_to_wandb function. """ """Log a table of validation samples with each example in its own row.""" if generations_to_log == 0: return Loading @@ -31,7 +28,7 @@ class ValidationTableManager: logger.warning('`val_generations_to_log_to_wandb` is set, but wandb is not initialized') return # Extract data from results exactly like training code # Extract data from results inputs = [] outputs = [] scores = [] Loading @@ -43,79 +40,67 @@ class ValidationTableManager: scores.append(item['metrics']['score']) images.append(item['image_data']) # Handle the case where images might not be provided if images is None or all(img is None for img in images): samples = list(zip(inputs, outputs, scores)) has_images = False max_images_per_sample = 0 else: samples = list(zip(inputs, outputs, scores, images)) has_images = True # Check if we have images has_images = any(img_list for img_list in images) # Find maximum number of images in any sample if has_images: max_images_per_sample = max( len(img_list) if isinstance(img_list, (list, tuple)) else 1 len(img_list) if img_list else 0 for img_list in images ) else: max_images_per_sample = 0 # Sort and shuffle exactly like training samples.sort(key=lambda x: x[0]) # Sort by input text # Create samples if has_images: samples = list(zip(inputs, outputs, scores, images)) else: samples = list(zip(inputs, outputs, scores)) # Use fixed random seed for deterministic shuffling rng = np.random.RandomState() # No seed specified in training code # Sort and shuffle for consistency samples.sort(key=lambda x: x[0]) # Sort by input text rng = np.random.RandomState() rng.shuffle(samples) # Take first N samples after shuffling # Take first N samples samples = samples[:generations_to_log] # Create column names for all samples # Create columns for the table if has_images: columns = ["step"] for i in range(len(samples)): columns.extend([f"input_{i+1}", f"output_{i+1}", f"score_{i+1}"]) columns.extend([f"image_{i+1}_{j+1}" for j in range(max_images_per_sample)]) columns = ["step", "input", "output", "score"] + [f"image_{i+1}" for i in range(max_images_per_sample)] else: columns = ["step"] + sum([[f"input_{i+1}", f"output_{i+1}", f"score_{i+1}"] for i in range(len(samples))], []) if self.validation_table is None: # Initialize the table on first call self.validation_table = wandb.Table(columns=columns) # Create a new table with same columns and existing data new_table = wandb.Table(columns=columns, data=self.validation_table.data) columns = ["step", "input", "output", "score"] # Add new row with all data row_data = [] row_data.append(global_steps) # Create table table = wandb.Table(columns=columns) # Add each sample as a separate row for sample in samples: if has_images: input_text, output_text, score, sample_images = sample row_data.extend([input_text, output_text, score]) # Handle if sample_images is a single image or list of images if not isinstance(sample_images, (list, tuple)): sample_images = [sample_images] # Convert each image to wandb.Image # Convert images to wandb.Image wandb_images = [] if sample_images: for img in sample_images: if img is not None: if not isinstance(img, wandb.Image): img = wandb.Image(img) wandb_images.append(img) # Pad with None if there are fewer images than max_images_per_sample wandb_images.extend([None] * (max_images_per_sample - len(wandb_images))) row_data.extend(wandb_images) # Pad with None if fewer images than max while len(wandb_images) < max_images_per_sample: wandb_images.append(None) # Add row table.add_data(global_steps, input_text, output_text, score, *wandb_images) else: input_text, output_text, score = sample row_data.extend([input_text, output_text, score]) new_table.add_data(*row_data) table.add_data(global_steps, input_text, output_text, score) # Update reference and log wandb.log({"val/generations": new_table}) self.validation_table = new_table # Log the table wandb.log({"val/generations": table}) # Global instance for maintaining table state across calls Loading @@ -141,29 +126,21 @@ def maybe_log_val_generations_to_wandb( def log_metrics_by_config_id( results: List[Dict[str, Any]], mode: str = 'inference' ) -> Dict[str, float]: ) -> Dict[str, Any]: """ Convert results to metrics dictionary with proper prefixes. Matches training code's log_rst_to_metrics_dict function. Convert results to metrics dictionary without any aggregation. Just logs raw metrics as they are. """ from collections import defaultdict metric_dict = {} # Group metrics by config_id metrics_by_config_id = defaultdict(dict) # dict of dict of list for item in results: # Simply iterate through results and create metric entries for i, item in enumerate(results): config_id = item["config_id"] env_id = item.get("env_id", f"env_{i}") # Log each metric as is, with a unique key including env_id for k, v in item["metrics"].items(): if k not in metrics_by_config_id[config_id]: metrics_by_config_id[config_id][k] = [] metrics_by_config_id[config_id][k].append(v) # Aggregate metrics for config_id, metrics in metrics_by_config_id.items(): for k, v in metrics.items(): metric_key = f'{mode}/{k}/{config_id}' metric_dict[metric_key] = np.mean(v) metric_key = f'{mode}/{config_id}/{env_id}/{k}' metric_dict[metric_key] = v return metric_dict No newline at end of file
vagen/mllm_agent/model_interface/__init__.py +5 −0 Original line number Diff line number Diff line from .vllm import VLLMModelInterface, VLLMModelConfig from .openai import OpenAIModelInterface, OpenAIModelConfig REGISTERED_MODEL = { "vllm": { "model_cls": VLLMModelInterface, "config_cls": VLLMModelConfig, }, "openai": { "model_cls": OpenAIModelInterface, "config_cls": OpenAIModelConfig } } No newline at end of file
vagen/mllm_agent/model_interface/openai/__init__.py 0 → 100644 +2 −0 Original line number Diff line number Diff line from .model import OpenAIModelInterface from .model_config import OpenAIModelConfig No newline at end of file
vagen/mllm_agent/model_interface/openai/model.py 0 → 100644 +221 −0 Original line number Diff line number Diff line # vagen/mllm_agent/model_interface/openai/model.py import base64 import logging import re from typing import List, Dict, Any from concurrent.futures import ThreadPoolExecutor from openai import OpenAI from PIL import Image import io from vagen.mllm_agent.model_interface.base_model import BaseModelInterface from .model_config import OpenAIModelConfig logger = logging.getLogger(__name__) class OpenAIModelInterface(BaseModelInterface): """Model interface for OpenAI API with Qwen format compatibility.""" def __init__(self, config: OpenAIModelConfig): super().__init__(config) self.config = config # Initialize OpenAI client self.client = OpenAI( api_key=config.api_key, organization=config.organization, base_url=config.base_url ) # Thread pool for batch processing self.executor = ThreadPoolExecutor(max_workers=10) logger.info(f"Initialized OpenAI interface with model {config.model_name}") def generate(self, prompts: List[Any], **kwargs) -> List[Dict[str, Any]]: """Generate responses using OpenAI API.""" # Process prompts into OpenAI message format formatted_requests = [] for prompt in prompts: messages = self._convert_qwen_to_openai_format(prompt) formatted_requests.append(messages) # Make parallel API calls futures = [] for messages in formatted_requests: future = self.executor.submit( self._single_api_call, messages, **kwargs ) futures.append(future) # Collect results results = [] for future in futures: try: result = future.result() results.append(result) except Exception as e: logger.error(f"API call failed: {e}") results.append({ "text": f"Error: {str(e)}", "error": str(e) }) return results def _convert_qwen_to_openai_format(self, prompt: List[Dict]) -> List[Dict]: """ Convert Qwen format messages to OpenAI format. Qwen format: Text with <image> placeholders + separate multi_modal_data OpenAI format: Structured content array with text and image objects """ openai_messages = [] for message in prompt: role = message.get("role", "user") content = message.get("content", "") # Create OpenAI message structure openai_msg = { "role": role, "content": [] } # Handle multimodal content if "multi_modal_data" in message and "<image>" in content: # Extract images from multi_modal_data images = [] for key, values in message["multi_modal_data"].items(): if key == "<image>" or "image" in key.lower(): images.extend(values) # Split content by <image> placeholders parts = content.split("<image>") # Build content array alternating text and images for i, part in enumerate(parts): # Add text part if not empty if part.strip(): openai_msg["content"].append({ "type": "text", "text": part }) # Add image if available (except for last part) if i < len(parts) - 1 and i < len(images): image_data = self._process_image_for_openai(images[i]) openai_msg["content"].append({ "type": "image_url", "image_url": { "url": f"data:image/jpeg;base64,{image_data}" } }) else: # Text-only message openai_msg["content"].append({ "type": "text", "text": content }) openai_messages.append(openai_msg) return openai_messages def _process_image_for_openai(self, image: Any) -> str: """Convert image to base64 for OpenAI API.""" if isinstance(image, Image.Image): # Ensure RGB mode if image.mode != "RGB": image = image.convert("RGB") # Resize if too large to save tokens max_size = 1024 if max(image.size) > max_size: ratio = max_size / max(image.size) new_size = tuple(int(dim * ratio) for dim in image.size) image = image.resize(new_size, Image.Resampling.LANCZOS) buffered = io.BytesIO() image.save(buffered, format="JPEG", quality=85) return base64.b64encode(buffered.getvalue()).decode() elif isinstance(image, dict) and "__pil_image__" in image: from vagen.server.serial import deserialize_pil_image pil_image = deserialize_pil_image(image) return self._process_image_for_openai(pil_image) else: raise ValueError(f"Unsupported image type: {type(image)}") def _single_api_call(self, messages: List[Dict], **kwargs) -> Dict[str, Any]: """Make a single API call to OpenAI.""" try: response = self.client.chat.completions.create( model=self.config.model_name, messages=messages, max_tokens=kwargs.get("max_tokens", self.config.max_tokens), temperature=kwargs.get("temperature", self.config.temperature), presence_penalty=kwargs.get("presence_penalty", self.config.presence_penalty), frequency_penalty=kwargs.get("frequency_penalty", self.config.frequency_penalty), seed=kwargs.get("seed", self.config.seed) ) # Extract text response response_text = response.choices[0].message.content # Convert response back to Qwen format if needed # (OpenAI doesn't return images, so no conversion needed for output) return { "text": response_text, "usage": { "prompt_tokens": response.usage.prompt_tokens, "completion_tokens": response.usage.completion_tokens, "total_tokens": response.usage.total_tokens }, "finish_reason": response.choices[0].finish_reason } except Exception as e: logger.error(f"OpenAI API error: {e}") raise def format_prompt(self, messages: List[Dict[str, Any]]) -> str: """ Format prompt for compatibility. Since OpenAI uses structured messages, this returns a string representation of the messages for logging/debugging purposes. """ formatted = [] for msg in messages: role = msg.get("role", "user") content = msg.get("content", "") # Handle Qwen special tokens if present if role == "system": formatted.append(f"System: {content}") elif role == "user": formatted.append(f"User: {content}") elif role == "assistant": formatted.append(f"Assistant: {content}") return "\n".join(formatted) def get_model_info(self) -> Dict[str, Any]: """Get detailed information about the model.""" info = super().get_model_info() info.update({ "name": self.config.model_name, "type": "multimodal" if "vision" in self.config.model_name.lower() else "text", "supports_images": "vision" in self.config.model_name.lower() or "4o" in self.config.model_name, "max_tokens": self.config.max_tokens, "temperature": self.config.temperature, "config_id": self.config.config_id() }) return info No newline at end of file
