langchain_huggingface

backend/models.py

@@ -1,11 +1,14 @@
 from typing import Any, Dict
 from enum import Enum
 
-from langchain_community.chat_models.huggingface import ChatHuggingFace
-from langchain_community.llms.huggingface_endpoint import HuggingFaceEndpoint
+#from langchain_community.chat_models.huggingface import ChatHuggingFace
+#from langchain_community.llms.huggingface_endpoint import HuggingFaceEndpoint
 from langchain_core import pydantic_v1
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.utils import get_from_dict_or_env
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
+
+
 from langchain_openai import ChatOpenAI
 
 
@@ -80,24 +83,33 @@ def get_chat_model_wrapper(
         model_id: str,
         inference_server_url: str,
         token: str,
-        backend: str = "HuggingFaceEndpoint",
+        backend: str = LLMBackends.HFChat,
         **model_init_kwargs
     ):
 
     backend = LLMBackends(backend)
 
     if backend == LLMBackends.HFChat:
-        llm = LazyHuggingFaceEndpoint(
-            endpoint_url=inference_server_url,
+        # llm = LazyHuggingFaceEndpoint(
+        #     endpoint_url=inference_server_url,
+        #     task="text-generation",
+        #     huggingfacehub_api_token=token,
+        #     **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(
+            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(
             model=model_id,

requirements.txt

@@ -2,8 +2,8 @@ gradio==4.37.2
 aiohttp
 datasets
 huggingface_hub
-langchain==0.2.4
-langchain_community==0.2.4
+langchain
+langchain_huggingface
 langchain_openai
 sentencepiece
 transformers