local-llm-kit 0.1.0__py3-none-any.whl

examples/__init__.py

@@ -0,0 +1,3 @@
+"""
+Example scripts for local_llm_kit.
+""" 

examples/simple_chat.py

@@ -0,0 +1,66 @@
+"""
+Simple example script for local_llm_kit chat.
+"""
+import argparse
+import sys
+
+from local_llm_kit import chat
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Simple chat with a local LLM")
+    parser.add_argument("--model", "-m", required=True, help="Path to the model or model name")
+    parser.add_argument("--backend", "-b", choices=["transformers", "llamacpp"], 
+                        help="Backend to use (default: auto-detect)")
+    args = parser.parse_args()
+    
+    # Initial message
+    print(f"Simple chat with {args.model}")
+    print("Type 'exit' or 'quit' to end the conversation.")
+    print()
+    
+    # Chat loop
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."}
+    ]
+    
+    try:
+        while True:
+            # Get user input
+            user_input = input("You: ")
+            if user_input.lower() in ["exit", "quit"]:
+                print("Goodbye!")
+                break
+                
+            # Add user message
+            messages.append({"role": "user", "content": user_input})
+            
+            # Generate and display response
+            try:
+                print("Assistant: ", end="", flush=True)
+                response_text = ""
+                
+                # Stream the response token by token
+                for chunk in chat(
+                    messages=messages,
+                    model_path=args.model,
+                    backend=args.backend,
+                    stream=True
+                ):
+                    # Extract and print chunk text
+                    if "choices" in chunk and chunk["choices"] and "delta" in chunk["choices"][0]:
+                        delta = chunk["choices"][0]["delta"]
+                        if "content" in delta and delta["content"]:
+                            print(delta["content"], end="", flush=True)
+                            response_text += delta["content"]
+                
+                print()  # Add newline after response
+                
+                # Add assistant response to history
+                messages.append({"role": "assistant", "content": response_text})
+                
+            except Exception as e:
+                print(f"\nError: {e}")
+                
+    except KeyboardInterrupt:
+        print("\nGoodbye!")
+        sys.exit(0) 

local_llm_kit/__init__.py

@@ -0,0 +1,11 @@
+"""
+local_llm_kit - OpenAI-like interface for local LLMs
+"""
+
+__version__ = "0.1.0"
+
+from .llm import LLM
+from .chat import chat, complete
+from .function_calling import add_function
+
+__all__ = ["LLM", "chat", "complete", "add_function"] 

local_llm_kit/backends/__init__.py

@@ -0,0 +1,9 @@
+"""
+Model backends for inference.
+"""
+
+from .base import BaseBackend
+from .transformers import TransformersBackend
+from .llamacpp import LlamaCppBackend
+
+__all__ = ["BaseBackend", "TransformersBackend", "LlamaCppBackend"] 

local_llm_kit/backends/base.py

@@ -0,0 +1,113 @@
+"""
+Base class for model backends.
+"""
+import time
+from abc import ABC, abstractmethod
+from typing import Dict, Any, Iterator, Optional
+
+
+class BaseBackend(ABC):
+    """
+    Abstract base class for model backends.
+    
+    All backend implementations must inherit from this class and
+    implement the required methods.
+    """
+    
+    @abstractmethod
+    def generate(
+        self, 
+        prompt: str, 
+        temperature: float = 0.7,
+        max_new_tokens: int = 512,
+        top_p: float = 0.95,
+        top_k: int = 40,
+        repetition_penalty: float = 1.1,
+        stream: bool = False,
+        logprobs: bool = False,
+        top_logprobs: Optional[int] = None,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Generate text completion for a prompt.
+        
+        Args:
+            prompt: Input text to complete
+            temperature: Sampling temperature (higher = more random)
+            max_new_tokens: Maximum number of tokens to generate
+            top_p: Top-p sampling parameter
+            top_k: Top-k sampling parameter
+            repetition_penalty: Penalty for repetition
+            stream: Whether to stream the response (should be False for this method)
+            logprobs: Whether to return log probabilities
+            top_logprobs: Number of top tokens to return logprobs for
+            **kwargs: Additional backend-specific parameters
+            
+        Returns:
+            Dictionary with at least a "text" key containing the generated text
+        """
+        pass
+    
+    @abstractmethod
+    def generate_stream(
+        self, 
+        prompt: str,
+        temperature: float = 0.7,
+        max_new_tokens: int = 512,
+        top_p: float = 0.95,
+        top_k: int = 40,
+        repetition_penalty: float = 1.1,
+        logprobs: bool = False,
+        top_logprobs: Optional[int] = None,
+        **kwargs
+    ) -> Iterator[Dict[str, Any]]:
+        """
+        Stream text completion for a prompt.
+        
+        Args:
+            prompt: Input text to complete
+            temperature: Sampling temperature (higher = more random)
+            max_new_tokens: Maximum number of tokens to generate
+            top_p: Top-p sampling parameter
+            top_k: Top-k sampling parameter
+            repetition_penalty: Penalty for repetition
+            logprobs: Whether to return log probabilities
+            top_logprobs: Number of top tokens to return logprobs for
+            **kwargs: Additional backend-specific parameters
+            
+        Yields:
+            Dictionaries with at least a "text" key containing a chunk of generated text
+        """
+        pass
+    
+    @abstractmethod
+    def get_context_window(self) -> int:
+        """
+        Get the context window size of the model.
+        
+        Returns:
+            Maximum context window size in tokens
+        """
+        pass
+    
+    @abstractmethod
+    def count_tokens(self, text: str) -> int:
+        """
+        Count the number of tokens in a text.
+        
+        Args:
+            text: Input text
+            
+        Returns:
+            Token count
+        """
+        pass
+    
+    def get_timestamp(self) -> int:
+        """
+        Get current timestamp.
+        
+        Returns:
+            Current timestamp in seconds
+        """
+        return int(time.time()) 

local_llm_kit/backends/llamacpp.py

@@ -0,0 +1,293 @@
+"""
+Backend for llama.cpp models via llama-cpp-python.
+"""
+import logging
+import time
+import json
+from typing import Dict, Any, Iterator, Optional, List, Callable, Union, Tuple
+
+from .base import BaseBackend
+
+logger = logging.getLogger(__name__)
+
+try:
+    from llama_cpp import Llama, LlamaGrammar
+    LLAMA_CPP_AVAILABLE = True
+except ImportError:
+    logger.warning("llama-cpp-python package not available. Install with 'pip install llama-cpp-python'")
+    LLAMA_CPP_AVAILABLE = False
+
+
+class LlamaCppBackend(BaseBackend):
+    """
+    Backend for llama.cpp models via llama-cpp-python binding.
+    """
+    
+    def __init__(
+        self, 
+        model_path: str,
+        n_gpu_layers: int = -1,
+        n_ctx: int = 2048,
+        n_batch: int = 512,
+        verbose: bool = False,
+        **kwargs
+    ):
+        """
+        Initialize the llama.cpp backend.
+        
+        Args:
+            model_path: Path to the model file (.gguf)
+            n_gpu_layers: Number of layers to offload to GPU (-1 = all)
+            n_ctx: Context window size
+            n_batch: Batch size for prompt evaluation
+            verbose: Whether to enable verbose output
+            **kwargs: Additional arguments to pass to Llama constructor
+        """
+        if not LLAMA_CPP_AVAILABLE:
+            raise ImportError(
+                "llama-cpp-python package is not installed. "
+                "Install it with 'pip install llama-cpp-python'"
+            )
+        
+        self.model_path = model_path
+        self.n_ctx = n_ctx
+        
+        # Load model
+        logger.info(f"Loading model from {model_path}")
+        self.model = Llama(
+            model_path=model_path,
+            n_gpu_layers=n_gpu_layers,
+            n_ctx=n_ctx,
+            n_batch=n_batch,
+            verbose=verbose,
+            **kwargs
+        )
+    
+    def generate(
+        self, 
+        prompt: str, 
+        temperature: float = 0.7,
+        max_new_tokens: int = 512,
+        top_p: float = 0.95,
+        top_k: int = 40,
+        repetition_penalty: float = 1.1,
+        stream: bool = False,
+        logprobs: bool = False,
+        top_logprobs: Optional[int] = None,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """Generate text completion for a prompt."""
+        if stream:
+            # If stream is requested, use generate_stream and collect all chunks
+            chunks = list(self.generate_stream(
+                prompt=prompt,
+                temperature=temperature,
+                max_new_tokens=max_new_tokens,
+                top_p=top_p,
+                top_k=top_k,
+                repetition_penalty=repetition_penalty,
+                logprobs=logprobs,
+                top_logprobs=top_logprobs,
+                **kwargs
+            ))
+            
+            # Combine chunks
+            text = "".join(chunk["text"] for chunk in chunks)
+            
+            # Combine logprobs if requested
+            combined_logprobs = None
+            if logprobs and chunks and "logprobs" in chunks[0]:
+                combined_logprobs = {
+                    "tokens": [],
+                    "token_logprobs": [],
+                    "top_logprobs": []
+                }
+                
+                for chunk in chunks:
+                    if "logprobs" in chunk:
+                        chunk_logprobs = chunk["logprobs"]
+                        combined_logprobs["tokens"].extend(chunk_logprobs.get("tokens", []))
+                        combined_logprobs["token_logprobs"].extend(chunk_logprobs.get("token_logprobs", []))
+                        
+                        if "top_logprobs" in chunk_logprobs:
+                            combined_logprobs["top_logprobs"].extend(chunk_logprobs.get("top_logprobs", []))
+            
+            result = {"text": text}
+            if combined_logprobs:
+                result["logprobs"] = combined_logprobs
+                
+            return result
+        
+        # Setup generation parameters
+        params = {
+            "temperature": temperature,
+            "max_tokens": max_new_tokens,
+            "top_p": top_p,
+            "top_k": top_k,
+            "repeat_penalty": repetition_penalty,
+            "logprobs": top_logprobs if logprobs else None,
+        }
+        
+        # Add grammar if provided
+        if "grammar" in kwargs:
+            grammar = kwargs.pop("grammar")
+            if isinstance(grammar, str):
+                params["grammar"] = LlamaGrammar.from_string(grammar)
+            elif isinstance(grammar, dict):
+                params["grammar"] = LlamaGrammar.from_string(json.dumps(grammar))
+        
+        # Generate
+        output = self.model(
+            prompt=prompt,
+            **params,
+            **kwargs
+        )
+        
+        # Extract text and tokens
+        if isinstance(output, dict):
+            text = output.get("choices", [{}])[0].get("text", "")
+        else:
+            # Handle list output format
+            text = output[0]["choices"][0]["text"]
+            
+        result = {"text": text}
+        
+        # Add logprobs if requested
+        if logprobs and top_logprobs:
+            try:
+                tokens_with_logprobs = []
+                
+                if isinstance(output, dict) and "choices" in output:
+                    choice = output["choices"][0]
+                    if "logprobs" in choice:
+                        tokens_with_logprobs = choice["logprobs"]["tokens"]
+                        token_logprobs = choice["logprobs"]["token_logprobs"]
+                        top_logprobs_data = choice["logprobs"]["top_logprobs"]
+                elif isinstance(output, list):
+                    choice = output[0]["choices"][0]
+                    if "logprobs" in choice:
+                        tokens_with_logprobs = choice["logprobs"]["tokens"]
+                        token_logprobs = choice["logprobs"]["token_logprobs"]
+                        top_logprobs_data = choice["logprobs"]["top_logprobs"]
+                
+                if tokens_with_logprobs:
+                    result["logprobs"] = {
+                        "tokens": tokens_with_logprobs,
+                        "token_logprobs": token_logprobs,
+                        "top_logprobs": top_logprobs_data
+                    }
+            except (KeyError, IndexError, TypeError) as e:
+                logger.warning(f"Failed to extract logprobs: {e}")
+            
+        return result
+    
+    def generate_stream(
+        self, 
+        prompt: str,
+        temperature: float = 0.7,
+        max_new_tokens: int = 512,
+        top_p: float = 0.95,
+        top_k: int = 40,
+        repetition_penalty: float = 1.1,
+        logprobs: bool = False,
+        top_logprobs: Optional[int] = None,
+        **kwargs
+    ) -> Iterator[Dict[str, Any]]:
+        """Stream text completion for a prompt."""
+        # Setup generation parameters
+        params = {
+            "temperature": temperature,
+            "max_tokens": max_new_tokens,
+            "top_p": top_p,
+            "top_k": top_k,
+            "repeat_penalty": repetition_penalty,
+            "logprobs": top_logprobs if logprobs else None,
+            "stream": True,
+        }
+        
+        # Add grammar if provided
+        if "grammar" in kwargs:
+            grammar = kwargs.pop("grammar")
+            if isinstance(grammar, str):
+                params["grammar"] = LlamaGrammar.from_string(grammar)
+            elif isinstance(grammar, dict):
+                params["grammar"] = LlamaGrammar.from_string(json.dumps(grammar))
+        
+        # Generate with streaming
+        for chunk in self.model(
+            prompt=prompt,
+            **params,
+            **kwargs
+        ):
+            # Extract delta text from chunk
+            if "choices" in chunk and chunk["choices"]:
+                delta_text = chunk["choices"][0].get("text", "")
+                
+                # Build chunk result
+                result = {"text": delta_text}
+                
+                # Add logprobs if available
+                if logprobs and top_logprobs:
+                    try:
+                        if "logprobs" in chunk["choices"][0]:
+                            logprobs_data = chunk["choices"][0]["logprobs"]
+                            result["logprobs"] = {
+                                "tokens": logprobs_data.get("tokens", []),
+                                "token_logprobs": logprobs_data.get("token_logprobs", []),
+                                "top_logprobs": logprobs_data.get("top_logprobs", [])
+                            }
+                    except (KeyError, IndexError, TypeError) as e:
+                        pass
+                
+                yield result
+    
+    def get_context_window(self) -> int:
+        """Get the context window size of the model."""
+        return self.n_ctx
+    
+    def count_tokens(self, text: str) -> int:
+        """Count the number of tokens in a text."""
+        return len(self.model.tokenize(text.encode("utf-8")))
+
+
+# Function to create JSON grammar for structured outputs
+def create_json_grammar() -> str:
+    """
+    Create a grammar for JSON generation.
+    
+    Returns:
+        Grammar string in BNF format
+    """
+    grammar = """
+        root ::= object
+        
+        value ::= object | array | string | number | ("true" | "false" | "null")
+        
+        object ::= "{" ws (string_colon_value (comma_string_colon_value)*)? ws "}"
+        
+        string_colon_value ::= ws string ws ":" ws value
+        
+        comma_string_colon_value ::= ws "," ws string ws ":" ws value
+        
+        array ::= "[" ws (value (comma_value)*)? ws "]"
+        
+        comma_value ::= ws "," ws value
+        
+        string ::= "\"" (
+            [^"\\] |
+            "\\\\" |
+            "\\\"" |
+            "\\/" |
+            "\\b" |
+            "\\f" |
+            "\\n" |
+            "\\r" |
+            "\\t" |
+            "\\u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]
+        )* "\""
+        
+        number ::= "-"? ("0" | [1-9] [0-9]*) ("." [0-9]+)? ([eE] [+-]? [0-9]+)?
+        
+        ws ::= [ \\t\\n\\r]*
+    """
+    return grammar