diff --git a/application/llm/llama_cpp.py b/application/llm/llama_cpp.py index a93fa686..804c3c56 100644 --- a/application/llm/llama_cpp.py +++ b/application/llm/llama_cpp.py @@ -1,8 +1,10 @@ from application.llm.base import BaseLLM from application.core.settings import settings +import threading class LlamaSingleton: _instances = {} + _lock = threading.Lock() # Add a lock for thread synchronization @classmethod def get_instance(cls, llm_name): @@ -16,6 +18,12 @@ class LlamaSingleton: cls._instances[llm_name] = Llama(model_path=llm_name, n_ctx=2048) return cls._instances[llm_name] + @classmethod + def query_model(cls, llm, prompt, **kwargs): + with cls._lock: + return llm(prompt, **kwargs) + + class LlamaCpp(BaseLLM): def __init__( self, @@ -34,14 +42,14 @@ class LlamaCpp(BaseLLM): context = messages[0]["content"] user_question = messages[-1]["content"] prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n" - result = self.llama(prompt, max_tokens=150, echo=False) + result = LlamaSingleton.query_model(self.llama, prompt, max_tokens=150, echo=False) return result["choices"][0]["text"].split("### Answer \n")[-1] def _raw_gen_stream(self, baseself, model, messages, stream=True, **kwargs): context = messages[0]["content"] user_question = messages[-1]["content"] prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n" - result = self.llama(prompt, max_tokens=150, echo=False, stream=stream) + result = LlamaSingleton.query_model(self.llama, prompt, max_tokens=150, echo=False, stream=stream) for item in result: for choice in item["choices"]: yield choice["text"] \ No newline at end of file