scitex 2.0.0__py2.py3-none-any.whl

scitex/ai/genai/llama.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-11-05 21:11:08 (ywatanabe)"
+# File: ./scitex_repo/src/scitex/ai/_gen_ai/_Llama.py
+
+"""
+DEPRECATED: This module is deprecated. Please use llama_provider.py instead.
+The new provider-based architecture offers better modularity and maintainability.
+"""
+
+"""Imports"""
+import os
+import sys
+from typing import List, Optional
+import warnings
+
+import matplotlib.pyplot as plt
+import scitex
+
+try:
+    from llama import Dialog
+    from llama import Llama as _Llama
+except:
+    pass
+
+from .base_genai import BaseGenAI
+
+warnings.warn(
+    "llama.py is deprecated. Please use llama_provider.py instead. "
+    "See PROVIDER_MIGRATION_GUIDE.md for migration instructions.",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+"""Functions & Classes"""
+
+
+def print_envs():
+    settings = {
+        "MASTER_ADDR": os.getenv("MASTER_ADDR", "localhost"),
+        "MASTER_PORT": os.getenv("MASTER_PORT", "12355"),
+        "WORLD_SIZE": os.getenv("WORLD_SIZE", "1"),
+        "RANK": os.getenv("RANK", "0"),
+    }
+
+    print("Environment Variable Settings:")
+    for key, value in settings.items():
+        print(f"{key}: {value}")
+    print()
+
+
+class Llama(BaseGenAI):
+    def __init__(
+        self,
+        ckpt_dir: str = "",
+        tokenizer_path: str = "",
+        system_setting: str = "",
+        model: str = "Meta-Llama-3-8B",
+        max_seq_len: int = 32_768,
+        max_batch_size: int = 4,
+        max_gen_len: Optional[int] = None,
+        stream: bool = False,
+        seed: Optional[int] = None,
+        n_keep: int = 1,
+        temperature: float = 1.0,
+        provider="Llama",
+        chat_history=None,
+        **kwargs,
+    ):
+
+        # Configure environment variables
+        os.environ["MASTER_ADDR"] = os.getenv("MASTER_ADDR", "localhost")
+        os.environ["MASTER_PORT"] = os.getenv("MASTER_PORT", "12355")
+        os.environ["WORLD_SIZE"] = os.getenv("WORLD_SIZE", "1")
+        os.environ["RANK"] = os.getenv("RANK", "0")
+        print_envs()
+
+        self.ckpt_dir = ckpt_dir if ckpt_dir else f"Meta-{model}/"
+        self.tokenizer_path = (
+            tokenizer_path if tokenizer_path else f"./Meta-{model}/tokenizer.model"
+        )
+        self.max_seq_len = max_seq_len
+        self.max_batch_size = max_batch_size
+        self.max_gen_len = max_gen_len
+
+        super().__init__(
+            system_setting=system_setting,
+            model=model,
+            api_key="",
+            stream=stream,
+            seed=seed,
+            n_keep=n_keep,
+            temperature=temperature,
+            chat_history=chat_history,
+        )
+
+    def __str__(self):
+        return "Llama"
+
+    def _init_client(self):
+        generator = _Llama.build(
+            ckpt_dir=self.ckpt_dir,
+            tokenizer_path=self.tokenizer_path,
+            max_seq_len=self.max_seq_len,
+            max_batch_size=self.max_batch_size,
+        )
+        return generator
+
+    def _api_call_static(self):
+        dialogs: List[Dialog] = [self.history]
+        results = self.client.chat_completion(
+            dialogs,
+            max_gen_len=self.max_gen_len,
+            temperature=self.temperature,
+            top_p=0.9,
+        )
+        out_text = results[0]["generation"]["content"]
+        return out_text
+
+    def _api_call_stream(self):
+        # Llama3 doesn't have built-in streaming, so we'll simulate it
+        full_response = self._api_call_static()
+        for char in full_response:
+            yield char
+
+    # def _get_available_models(self):
+    #     # Llama3 doesn't have a list of available models, so we'll return a placeholder
+    #     return ["llama3"]
+
+    def verify_model(self):
+        # Llama3 doesn't require model verification, so we'll skip it
+        pass
+
+
+def main():
+    m = Llama(
+        ckpt_dir="/path/to/checkpoint",
+        tokenizer_path="/path/to/tokenizer",
+        system_setting="You are a helpful assistant.",
+        max_seq_len=512,
+        max_batch_size=4,
+        stream=True,
+        temperature=0.7,
+    )
+    m("Hi")
+    pass
+
+
+if __name__ == "__main__":
+    # Main
+    CONFIG, sys.stdout, sys.stderr, plt, CC = scitex.gen.start(sys, plt, verbose=False)
+    main()
+    scitex.gen.close(CONFIG, verbose=False, notify=False)
+
+# EOF

scitex/ai/genai/llama_provider.py

@@ -0,0 +1,214 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-11-13 20:25:55 (ywatanabe)"
+# File: ./scitex_repo/src/scitex/ai/genai/llama_provider.py
+
+"""Llama provider implementation using the new component-based architecture.
+
+This module provides integration with local Llama models through the official Llama library.
+It supports loading and running Llama models locally with full control over model parameters.
+"""
+
+import os
+from typing import Dict, List, Iterator, Optional, Any
+
+try:
+    from llama import Llama as _Llama
+    from llama import Dialog
+except ImportError:
+    _Llama = None
+    Dialog = None
+    print(
+        "Warning: llama package not installed. "
+        "Install with the official Meta Llama repository instructions."
+    )
+
+from .base_provider import BaseProvider, ProviderConfig
+
+
+class LlamaProvider(BaseProvider):
+    """Llama provider implementation for local model inference."""
+
+    SUPPORTED_MODELS = [
+        "Meta-Llama-3-8B",
+        "Meta-Llama-3-70B",
+        "Meta-Llama-3.1-8B",
+        "Meta-Llama-3.1-70B",
+        "Meta-Llama-3.1-405B",
+        "Llama-2-7b",
+        "Llama-2-13b",
+        "Llama-2-70b",
+    ]
+
+    DEFAULT_MODEL = "Meta-Llama-3-8B"
+
+    def __init__(self, config: ProviderConfig):
+        """Initialize Llama provider.
+
+        Args:
+            config: Provider configuration
+        """
+        self.config = config
+        self.model_name = config.model or self.DEFAULT_MODEL
+
+        # Llama-specific configuration
+        self.ckpt_dir = getattr(config, "ckpt_dir", None) or f"{self.model_name}/"
+        self.tokenizer_path = (
+            getattr(config, "tokenizer_path", None)
+            or f"{self.model_name}/tokenizer.model"
+        )
+        self.max_seq_len = getattr(config, "max_seq_len", 32_768)
+        self.max_batch_size = getattr(config, "max_batch_size", 4)
+        self.max_gen_len = config.max_tokens
+
+        # Configure environment variables for distributed inference
+        self._setup_environment()
+
+        # Initialize the Llama model
+        if _Llama is None:
+            raise ImportError(
+                "Llama package is not installed. Please install it from the official Meta repository."
+            )
+
+        try:
+            self.generator = _Llama.build(
+                ckpt_dir=self.ckpt_dir,
+                tokenizer_path=self.tokenizer_path,
+                max_seq_len=self.max_seq_len,
+                max_batch_size=self.max_batch_size,
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to load Llama model: {str(e)}")
+
+    def _setup_environment(self):
+        """Set up environment variables for distributed inference."""
+        env_vars = {
+            "MASTER_ADDR": os.getenv("MASTER_ADDR", "localhost"),
+            "MASTER_PORT": os.getenv("MASTER_PORT", "12355"),
+            "WORLD_SIZE": os.getenv("WORLD_SIZE", "1"),
+            "RANK": os.getenv("RANK", "0"),
+        }
+
+        for key, value in env_vars.items():
+            os.environ[key] = value
+
+    def validate_messages(self, messages: List[Dict[str, Any]]) -> None:
+        """Validate message format.
+
+        Args:
+            messages: List of message dictionaries
+
+        Raises:
+            ValueError: If messages are invalid
+        """
+        if not messages:
+            raise ValueError("Messages cannot be empty")
+
+        for msg in messages:
+            if "role" not in msg:
+                raise ValueError(f"Missing role in message: {msg}")
+            if "content" not in msg:
+                raise ValueError(f"Missing content in message: {msg}")
+            if msg["role"] not in ["system", "user", "assistant"]:
+                raise ValueError(f"Invalid role: {msg['role']}")
+
+    def format_messages(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Format messages for Llama API.
+
+        Args:
+            messages: List of message dictionaries
+
+        Returns:
+            Formatted messages for Llama
+        """
+        formatted = []
+
+        # Add system prompt if configured
+        if self.config.system_prompt:
+            formatted.append({"role": "system", "content": self.config.system_prompt})
+
+        # Add user messages
+        formatted.extend(messages)
+
+        return formatted
+
+    def complete(self, messages: List[Dict[str, Any]], **kwargs) -> Dict[str, Any]:
+        """Generate a completion.
+
+        Args:
+            messages: List of message dictionaries
+            **kwargs: Additional parameters for the API
+
+        Returns:
+            Completion response dictionary
+        """
+        self.validate_messages(messages)
+        formatted_messages = self.format_messages(messages)
+
+        # Convert to Llama Dialog format
+        dialogs: List[Dialog] = [formatted_messages]
+
+        # Merge config parameters with kwargs
+        params = {
+            "max_gen_len": self.max_gen_len,
+            "temperature": self.config.temperature or 1.0,
+            "top_p": kwargs.get("top_p", 0.9),
+        }
+
+        try:
+            results = self.generator.chat_completion(dialogs, **params)
+
+            result = results[0]
+            content = result["generation"]["content"]
+
+            # Estimate token counts (Llama doesn't provide exact counts)
+            prompt_tokens = len(
+                " ".join(msg["content"] for msg in formatted_messages).split()
+            )
+            completion_tokens = len(content.split())
+
+            return {
+                "content": content,
+                "model": self.model_name,
+                "usage": {
+                    "prompt_tokens": prompt_tokens,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": prompt_tokens + completion_tokens,
+                },
+                "finish_reason": "stop",
+            }
+        except Exception as e:
+            raise RuntimeError(f"Llama inference error: {str(e)}")
+
+    def stream(
+        self, messages: List[Dict[str, Any]], **kwargs
+    ) -> Iterator[Dict[str, Any]]:
+        """Stream a completion.
+
+        Note: Llama doesn't have native streaming support, so this simulates streaming
+        by yielding characters one at a time.
+
+        Args:
+            messages: List of message dictionaries
+            **kwargs: Additional parameters for the API
+
+        Yields:
+            Chunks of the completion
+        """
+        # Get the full response
+        response = self.complete(messages, **kwargs)
+        content = response["content"]
+
+        # Simulate streaming by yielding characters
+        for i, char in enumerate(content):
+            yield {
+                "content": char,
+                "model": self.model_name,
+            }
+
+        # Yield final chunk with usage info
+        yield {
+            "content": "",
+            "usage": response["usage"],
+            "finish_reason": "stop",
+        }

scitex/ai/genai/mock_provider.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-11-15 12:00:00"
+# Author: Yusuke Watanabe (ywatanabe@alumni.u-tokyo.ac.jp)
+
+"""
+Mock provider for testing purposes.
+"""
+
+from typing import List, Dict, Any, Optional, Iterator, Generator
+from .base_provider import BaseProvider, CompletionResponse, Provider
+from .provider_factory import register_provider
+
+
+class MockProvider(BaseProvider):
+    """Mock provider for testing."""
+
+    def __init__(self, config):
+        """Initialize mock provider."""
+        self.config = config
+        self.api_key = config.api_key
+        self.model = config.model or "mock-model"
+        self.stream_mode = config.stream
+        self.system_prompt = config.system_prompt
+        self.temperature = config.temperature
+        self.max_tokens = config.max_tokens
+        self.seed = config.seed
+        self.n_draft = config.n_draft
+        self.client = None  # Mock client
+
+    def init_client(self) -> Any:
+        """Initialize the mock client."""
+        self.client = {"mock": True}  # Mock client object
+        return self.client
+
+    def format_history(self, history: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Format conversation history."""
+        # Mock implementation - just return as-is
+        return history
+
+    def call_static(self, messages: List[Dict[str, Any]], **kwargs) -> Any:
+        """Make a static API call."""
+        # Mock response
+        content = f"Mock response to: {messages[-1]['content']}"
+        return {
+            "choices": [
+                {
+                    "message": {"content": content, "role": "assistant"},
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": {
+                "prompt_tokens": len(str(messages)),
+                "completion_tokens": len(content),
+                "total_tokens": len(str(messages)) + len(content),
+            },
+        }
+
+    def call_stream(
+        self, messages: List[Dict[str, Any]], **kwargs
+    ) -> Generator[str, None, None]:
+        """Make a streaming API call."""
+        response = f"Mock streaming response to: {messages[-1]['content']}"
+        for word in response.split():
+            yield word + " "
+
+    def complete(self, messages: List[Dict[str, Any]], **kwargs) -> CompletionResponse:
+        """Generate a mock completion."""
+        # Ensure client is initialized
+        if not self.client:
+            self.init_client()
+
+        # Format history
+        formatted_messages = self.format_history(messages)
+
+        # Make API call
+        response = self.call_static(formatted_messages, **kwargs)
+
+        # Extract content
+        content = response["choices"][0]["message"]["content"]
+        usage = response.get("usage", {})
+
+        return CompletionResponse(
+            content=content,
+            input_tokens=usage.get("prompt_tokens", 0),
+            output_tokens=usage.get("completion_tokens", 0),
+            finish_reason=response["choices"][0].get("finish_reason", "stop"),
+            provider_response=response,
+        )
+
+    def stream(self, messages: List[Dict[str, Any]], **kwargs) -> Iterator[str]:
+        """Stream mock completions."""
+        # Ensure client is initialized
+        if not self.client:
+            self.init_client()
+
+        # Format history
+        formatted_messages = self.format_history(messages)
+
+        # Stream response
+        for chunk in self.call_stream(formatted_messages, **kwargs):
+            yield chunk
+
+    def count_tokens(self, text: str) -> int:
+        """Count tokens in text."""
+        return len(text) // 4  # Mock implementation
+
+    @property
+    def supports_images(self) -> bool:
+        """Check if provider supports images."""
+        return True
+
+    @property
+    def supports_streaming(self) -> bool:
+        """Check if provider supports streaming."""
+        return True
+
+    @property
+    def max_context_length(self) -> int:
+        """Get maximum context length."""
+        return 4096
+
+
+# Auto-register when module is imported
+register_provider(Provider.MOCK.value, MockProvider)
+
+## EOF