LazyHfEndpoint validator

backend/models.py

@@ -46,6 +46,10 @@ class LazyHuggingFaceEndpoint(HuggingFaceEndpoint):
     # which might in fact be a hf_oauth token that does only permit inference,
     # not logging in.
 
+    @pydantic_v1.root_validator(pre=True)
+    def build_extra(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+        return super().build_extra(values)
+
     @pydantic_v1.root_validator()
     def validate_environment(cls, values: Dict) -> Dict:  # noqa: UP006, N805
         """Validate that package is installed and that the API token is valid."""
@@ -83,7 +87,7 @@ def get_chat_model_wrapper(
         model_id: str,
         inference_server_url: str,
         token: str,
-        backend: str = LLMBackends.HFChat,
+        backend: str = "HFChat",
         **model_init_kwargs
     ):
 
@@ -97,18 +101,17 @@ def get_chat_model_wrapper(
         #     **model_init_kwargs,
         # )
 
-        # from transformers import AutoTokenizer
-
-        # tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
-        # chat_model = LazyChatHuggingFace(llm=llm, model_id=model_id, tokenizer=tokenizer)
-
-        llm = HuggingFaceEndpoint(
+        llm = LazyHuggingFaceEndpoint(
             repo_id=model_id,
             task="text-generation",
             huggingfacehub_api_token=token,
             **model_init_kwargs,
         )
-        chat_model = ChatHuggingFace(llm=llm)
+
+        from transformers import AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
+        chat_model = LazyChatHuggingFace(llm=llm, model_id=model_id, tokenizer=tokenizer)
 
     elif backend in [LLMBackends.VLLM, LLMBackends.Fireworks]:
         chat_model = ChatOpenAI(
@@ -118,4 +121,7 @@ def get_chat_model_wrapper(
             **model_init_kwargs,
         )
 
+    else:
+        raise ValueError(f"Backend {backend} not supported")
+
     return chat_model

requirements.txt

@@ -2,9 +2,10 @@ gradio==4.37.2
 aiohttp
 datasets
 huggingface_hub
-langchain
-langchain_huggingface
-langchain_openai
+langchain==0.2.7
+langchain_core==0.2.16
+langchain_huggingface==0.0.3
+langchain-openai==0.1.15
 sentencepiece
 transformers
 ujson