Spaces:
Running
Running
File size: 6,036 Bytes
31086ae 0134a08 31086ae 0134a08 31086ae 0134a08 31086ae 982cb95 0134a08 31086ae 0134a08 31086ae 0134a08 31086ae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 |
import os
from typing import Any, Dict, Optional
import ray
from graphgen.bases import BaseLLMWrapper
from graphgen.models import Tokenizer
class LLMServiceActor:
"""
A Ray actor class to wrap LLM wrapper instances for distributed usage.
"""
def __init__(self, backend: str, config: Dict[str, Any]):
self.backend = backend
tokenizer_model = os.environ.get("TOKENIZER_MODEL", "cl100k_base")
tokenizer = Tokenizer(model_name=tokenizer_model)
config["tokenizer"] = tokenizer
if backend == "http_api":
from graphgen.models.llm.api.http_client import HTTPClient
self.llm_instance = HTTPClient(**config)
elif backend in ("openai_api", "azure_openai_api"):
from graphgen.models.llm.api.openai_client import OpenAIClient
# pass in concrete backend to the OpenAIClient so that internally we can distinguish
# between OpenAI and Azure OpenAI
self.llm_instance = OpenAIClient(**config, backend=backend)
elif backend == "ollama_api":
from graphgen.models.llm.api.ollama_client import OllamaClient
self.llm_instance = OllamaClient(**config)
elif backend == "huggingface":
from graphgen.models.llm.local.hf_wrapper import HuggingFaceWrapper
self.llm_instance = HuggingFaceWrapper(**config)
elif backend == "sglang":
from graphgen.models.llm.local.sglang_wrapper import SGLangWrapper
self.llm_instance = SGLangWrapper(**config)
elif backend == "vllm":
from graphgen.models.llm.local.vllm_wrapper import VLLMWrapper
self.llm_instance = VLLMWrapper(**config)
else:
raise NotImplementedError(f"Backend {backend} is not implemented yet.")
async def generate_answer(
self, text: str, history: Optional[list[str]] = None, **extra: Any
) -> str:
return await self.llm_instance.generate_answer(text, history, **extra)
async def generate_topk_per_token(
self, text: str, history: Optional[list[str]] = None, **extra: Any
) -> list:
return await self.llm_instance.generate_topk_per_token(text, history, **extra)
async def generate_inputs_prob(
self, text: str, history: Optional[list[str]] = None, **extra: Any
) -> list:
return await self.llm_instance.generate_inputs_prob(text, history, **extra)
def ready(self) -> bool:
"""A simple method to check if the actor is ready."""
return True
class LLMServiceProxy(BaseLLMWrapper):
"""
A proxy class to interact with the LLMServiceActor for distributed LLM operations.
"""
def __init__(self, actor_handle: ray.actor.ActorHandle):
super().__init__()
self.actor_handle = actor_handle
self._create_local_tokenizer()
async def generate_answer(
self, text: str, history: Optional[list[str]] = None, **extra: Any
) -> str:
object_ref = self.actor_handle.generate_answer.remote(text, history, **extra)
return await object_ref
async def generate_topk_per_token(
self, text: str, history: Optional[list[str]] = None, **extra: Any
) -> list:
object_ref = self.actor_handle.generate_topk_per_token.remote(
text, history, **extra
)
return await object_ref
async def generate_inputs_prob(
self, text: str, history: Optional[list[str]] = None, **extra: Any
) -> list:
object_ref = self.actor_handle.generate_inputs_prob.remote(
text, history, **extra
)
return await object_ref
def _create_local_tokenizer(self):
tokenizer_model = os.environ.get("TOKENIZER_MODEL", "cl100k_base")
self.tokenizer = Tokenizer(model_name=tokenizer_model)
class LLMFactory:
"""
A factory class to create LLM wrapper instances based on the specified backend.
Supported backends include:
- http_api: HTTPClient
- openai_api: OpenAIClient
- ollama_api: OllamaClient
- huggingface: HuggingFaceWrapper
- sglang: SGLangWrapper
"""
@staticmethod
def create_llm(
model_type: str, backend: str, config: Dict[str, Any]
) -> BaseLLMWrapper:
if not config:
raise ValueError(
f"No configuration provided for LLM {model_type} with backend {backend}."
)
actor_name = f"Actor_LLM_{model_type}"
try:
actor_handle = ray.get_actor(actor_name)
print(f"Using existing Ray actor: {actor_name}")
except ValueError:
print(f"Creating Ray actor for LLM {model_type} with backend {backend}.")
num_gpus = float(config.pop("num_gpus", 0))
actor_handle = (
ray.remote(LLMServiceActor)
.options(
name=actor_name,
num_gpus=num_gpus,
get_if_exists=True,
)
.remote(backend, config)
)
# wait for actor to be ready
ray.get(actor_handle.ready.remote())
return LLMServiceProxy(actor_handle)
def _load_env_group(prefix: str) -> Dict[str, Any]:
"""
Collect environment variables with the given prefix into a dictionary,
stripping the prefix from the keys.
"""
return {
k[len(prefix) :].lower(): v
for k, v in os.environ.items()
if k.startswith(prefix)
}
def init_llm(model_type: str) -> Optional[BaseLLMWrapper]:
if model_type == "synthesizer":
prefix = "SYNTHESIZER_"
elif model_type == "trainee":
prefix = "TRAINEE_"
else:
raise NotImplementedError(f"Model type {model_type} is not implemented yet.")
config = _load_env_group(prefix)
# if config is empty, return None
if not config:
return None
backend = config.pop("backend")
llm_wrapper = LLMFactory.create_llm(model_type, backend, config)
return llm_wrapper
|