my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

my-project-public — gitlawb

#1	from typing import Iterable, Optional, Union
#2
#3	from langchain.callbacks.manager import CallbackManager
#4	from langchain.callbacks.stdout import StdOutCallbackHandler
#5	from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
#6	from langchain_community.llms import VLLM as BaseVLLM
#7
#8	from embedchain.config import BaseLlmConfig
#9	from embedchain.helpers.json_serializable import register_deserializable
#10	from embedchain.llm.base import BaseLlm
#11
#12
#13	@register_deserializable
#14	class VLLM(BaseLlm):
#15	def __init__(self, config: Optional[BaseLlmConfig] = None):
#16	super().__init__(config=config)
#17	if self.config.model is None:
#18	self.config.model = "mosaicml/mpt-7b"
#19
#20	def get_llm_model_answer(self, prompt):
#21	return self._get_answer(prompt=prompt, config=self.config)
#22
#23	@staticmethod
#24	def _get_answer(prompt: str, config: BaseLlmConfig) -> Union[str, Iterable]:
#25	callback_manager = [StreamingStdOutCallbackHandler()] if config.stream else [StdOutCallbackHandler()]
#26
#27	# Prepare the arguments for BaseVLLM
#28	llm_args = {
#29	"model": config.model,
#30	"temperature": config.temperature,
#31	"top_p": config.top_p,
#32	"callback_manager": CallbackManager(callback_manager),
#33	}
#34
#35	# Add model_kwargs if they are not None
#36	if config.model_kwargs is not None:
#37	llm_args.update(config.model_kwargs)
#38
#39	llm = BaseVLLM(**llm_args)
#40	return llm.invoke(prompt)
#41

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public