llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Allow stopping of inference

Browse Source
This commit is contained in:
geens
2024-11-30 12:44:13 +01:00
parent a0353d0d49
commit e71bd7e9eb
15 changed files with 139 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
sia/__main__.py
View File

@@ -34,7 +34,8 @@ class Main:
self._llms['local'] = LocalLlmEngine(
config.local_model,
config.local_temperature,
config.local_token_limit
config.local_token_limit,
config.local_api_key,
)
if config.openai_enabled:

10
sia/config.py
View File

@@ -82,6 +82,12 @@ class Config:
default=int(os.getenv('SIA_LOCAL_TOKEN_LIMIT', '2048')),
help='Local LLM token limit (env: SIA_LOCAL_TOKEN_LIMIT)'
)
parser.add_argument(
'--local-api-key',
type=str,
default=os.getenv('SIA_LOCAL_API_KEY'),
help='API key for local models (env: SIA_LOCAL_API_KEY)'
)
# OpenAI configuration
parser.add_argument(
@@ -233,6 +239,10 @@ class Config:
@property
def local_token_limit(self) -> int:
return self.args.local_token_limit
@property
def local_api_key(self) -> Optional[str]:
return self.args.local_api_key
# OpenAI properties
@property

29
sia/hf_llm_engine.py
View File

@@ -1,6 +1,6 @@
from huggingface_hub import InferenceClient
from transformers import AutoTokenizer, AutoConfig
from typing import Iterator, Optional
from typing import Iterator, Optional, Callable
from .llm_engine import LlmEngine
@@ -30,36 +30,39 @@ class HfLlmEngine(LlmEngine):
self._config = AutoConfig.from_pretrained(model, token=api_token)
self._client = InferenceClient(token=api_token)
def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
"""
Run inference using the system prompt and main context.
Args:
system_prompt: The system prompt string
main_context: The main context string after templating
should_stop: Callback that returns True when inference should stop
Returns:
Iterator[str]: An iterator that yields the generated text.
"""
token_count=self.token_count(system_prompt, main_context)
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
]
def stream_wrapper():
stream = self._client.chat_completion(
model=self._model,
messages=messages,
temperature=self._temperature,
stream=True
)
stream = self._client.chat_completion(
model=self._model,
messages=messages,
temperature=self._temperature,
stream=True
)
try:
for response in stream:
if should_stop():
stream.close()
break
if content := response.choices[0].delta.content:
yield content
return stream_wrapper()
finally:
stream.close()
def token_count(self, system_prompt: str, main_context: str) -> int:
messages = [

4
sia/llm_engine.py
View File

@@ -1,9 +1,9 @@
from typing import Iterator
from typing import Callable, Iterator
from abc import ABC, abstractmethod
class LlmEngine(ABC):
@abstractmethod
def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
pass
@abstractmethod

24
sia/local_llm_engine.py
View File

@@ -1,5 +1,5 @@
from threading import Thread
from typing import Iterator, Optional
from typing import Iterator, Optional, Callable
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
import torch
@@ -49,13 +49,14 @@ class LocalLlmEngine(LlmEngine):
return_full_text=False,
)
def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
"""
Run inference using the system prompt and main context, while validating actions against the provided XML schema.
Run inference using the system prompt and main context.
Args:
system_prompt: The system prompt string
main_context: The main context string after templating
should_stop: Callback that returns True when inference should stop
Returns:
Iterator[str]: An iterator that yields the generated text.
@@ -71,16 +72,21 @@ class LocalLlmEngine(LlmEngine):
self._tokenizer,
skip_prompt=True
)
pipeline_kwargs = dict(
generation_thread = Thread(target=self._pipeline, kwargs=dict(
text_inputs=prompt,
do_sample=True,
temperature=self._temperature,
max_new_tokens=self.token_limit(),
streamer=streamer
)
thread = Thread(target=self._pipeline, kwargs=pipeline_kwargs)
thread.start()
return util.stop_before_value(streamer, '<|eot_id|>')
))
generation_thread.start()
for text in util.stop_before_value(streamer, '<|eot_id|>'):
yield text
if should_stop():
break
generation_thread.join()
def token_count(self, system_prompt: str, main_context: str) -> int:
"""
@@ -112,4 +118,4 @@ class LocalLlmEngine(LlmEngine):
if self._token_limit is not None:
return self._token_limit
else:
return self._pipeline.model.config.max_position_embeddings
return self._pipeline.model.config.max_position_embeddings

25
sia/mistral_llm_engine.py
View File

@@ -1,6 +1,4 @@
from typing import Iterator, Optional
from abc import ABC, abstractmethod
import os
from typing import Iterator, Optional, Callable
from mistralai import Mistral
from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage
from mistral_common.protocol.instruct.request import ChatCompletionRequest
@@ -23,11 +21,7 @@ class MistralLlmEngine(LlmEngine):
self._client = Mistral(api_key=api_key)
self._tokenizer = MistralTokenizer.v3()
def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
messages = [
SystemMessage(content=system_prompt),
UserMessage(content=main_context),
]
def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
messages = [
{
"role": "system",
@@ -48,9 +42,16 @@ class MistralLlmEngine(LlmEngine):
messages=messages,
temperature=self._temperature,
)
for chunk in stream_response:
if content := chunk.data.choices[0].delta.content:
yield chunk.data.choices[0].delta.content
try:
for chunk in stream_response:
if should_stop():
stream_response.close()
break
if content := chunk.data.choices[0].delta.content:
yield content
finally:
stream_response.close()
def token_count(self, system_prompt: str, main_context: str) -> int:
messages = [
@@ -67,4 +68,4 @@ class MistralLlmEngine(LlmEngine):
return len(tokenized.tokens)
def token_limit(self) -> int:
return self._token_limit
return self._token_limit

26
sia/openai_llm_engine.py
View File

@@ -1,4 +1,4 @@
from typing import Iterator
from typing import Callable, Iterator
import openai
import tiktoken
@@ -34,17 +34,7 @@ class OpenAILlmEngine(LlmEngine):
api_key=api_key,
)
def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
"""
Run inference using the system prompt and main context.
Args:
system_prompt: The system prompt string
main_context: The main context string after templating
Returns:
Iterator[str]: An iterator that yields the generated text in chunks.
"""
def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": main_context}
@@ -57,9 +47,15 @@ class OpenAILlmEngine(LlmEngine):
stream=True,
)
for chunk in stream:
if content := chunk.choices[0].delta.content:
yield content
try:
for chunk in stream:
if should_stop():
break
if content := chunk.choices[0].delta.content:
yield content
finally:
stream.close()
#stream.response.close()
def token_count(self, system_prompt: str, main_context: str) -> int:
"""

10
sia/web/api.py
View File

@@ -30,6 +30,7 @@ class Api:
def _init_routes(self):
"""Initialize REST API and WebSocket routes."""
self._app.router.add_post("/api/inference/{llm}", self._run_inference)
self._app.router.add_post("/api/inference/{llm}/stop", self._stop_inference)
self._app.router.add_post("/api/approve/{llm}", self._approve_response)
self._app.router.add_post("/api/context", self._modify_context)
self._app.router.add_post("/api/input", self._send_input)
@@ -60,6 +61,15 @@ class Api:
except (ValueError, RuntimeError) as e:
return web.Response(status=400, text=str(e))
async def _stop_inference(self, request: web.Request) -> web.Response:
"""Stop inference on specified LLM."""
llm_name = request.match_info["llm"]
try:
self._agent.stop_inference(llm_name)
return web.Response(status=200)
except ValueError as e:
return web.Response(status=400, text=str(e))
async def _approve_response(self, request: web.Request) -> web.Response:
"""Approve response from specified LLM."""
llm_name = request.match_info["llm"]

16
sia/web_agent.py
View File

@@ -46,6 +46,7 @@ class WebAgent(BaseAgent):
self._validation_error: Optional[str] = None
self._command_result: Optional[CommandResult] = None
self._context = self._compile_context(next(iter(self._llms.values())))
self._stop_flags: Dict[str, bool] = {name: False for name in llms}
# Locks
self._llm_lock = Lock()
@@ -112,9 +113,14 @@ class WebAgent(BaseAgent):
if self._llm_states[llm_name] != LlmState.NO_OUTPUT:
raise RuntimeError(f"LLM {llm_name} is not ready for inference")
self._set_llm_state(llm_name, LlmState.INFERENCE)
self._stop_flags[llm_name] = False
llm = self._llms[llm_name]
response_token_iter = llm.infer(self.system_prompt, self.context)
def should_stop() -> bool:
return self._stop_flags[llm_name]
response_token_iter = llm.infer(self.system_prompt, self.context, should_stop)
with self._output_lock:
self._llm_outputs[llm_name] = ""
@@ -129,6 +135,12 @@ class WebAgent(BaseAgent):
with self._llm_lock:
self._set_llm_state(llm_name, LlmState.OUTPUT)
def stop_inference(self, llm_name: str) -> None:
"""Stop ongoing inference for specified LLM"""
if llm_name not in self._llms:
raise ValueError(f"Unknown LLM: {llm_name}")
self._stop_flags[llm_name] = True
def get_output(self, llm_name: str) -> str:
"""Get complete output for specified LLM"""
if llm_name not in self._llms:
@@ -164,4 +176,4 @@ class WebAgent(BaseAgent):
def _handle_memory_update(self) -> None:
"""Handle memory updates and update context"""
context = self._compile_context(next(iter(self._llms.values())))
self.modify_context(context, True)
self.modify_context(context, True)