added support for HF TGI

.env.example

@@ -6,6 +6,7 @@ LLM_MODEL_TYPE=huggingface
 # LLM_MODEL_TYPE=mosaicml
 # LLM_MODEL_TYPE=stablelm
 # LLM_MODEL_TYPE=openllm
+# LLM_MODEL_TYPE=hftgi
 
 OPENLLM_SERVER_URL=http://localhost:64300
 

.gitignore

@@ -1,5 +1,7 @@
 pdfs
 .vscode/
+data/version.txt
+data/models*
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

app_modules/llm_loader.py

@@ -5,6 +5,7 @@ from queue import Queue
 from typing import Any, Optional
 
 import torch
+from langchain import HuggingFaceTextGenInference
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 from langchain.chat_models import ChatOpenAI
@@ -188,6 +189,19 @@ class LLMLoader:
                     verbose=True,
                     use_mlock=True,
                 )
+            elif self.llm_model_type == "hftgi":
+                HFTGI_SERVER_URL = os.environ.get("HFTGI_SERVER_URL")
+                self.llm = HuggingFaceTextGenInference(
+                    inference_server_url=HFTGI_SERVER_URL,
+                    max_new_tokens=self.max_tokens_limit / 2,
+                    top_k=10,
+                    top_p=0.95,
+                    typical_p=0.95,
+                    temperature=0.01,
+                    repetition_penalty=1.03,
+                    callbacks=callbacks,
+                    streaming=True,
+                )
             elif self.llm_model_type.startswith("huggingface"):
                 MODEL_NAME_OR_PATH = os.environ.get("HUGGINGFACE_MODEL_NAME_OR_PATH")
                 print(f"            loading model: {MODEL_NAME_OR_PATH}")

requirements.txt

@@ -34,4 +34,5 @@ pypdf
 python-telegram-bot
 transformers_stream_generator
 openllm
-openllm[llama]
+openllm[llama]
+text_generation

tgi.sh

@@ -0,0 +1,9 @@
+#!/bin/sh
+
+export HUGGINGFACE_MODEL_NAME_OR_PATH="meta-llama/Llama-2-7b-chat-hf"
+
+echo Running $HUGGINGFACE_MODEL_NAME_OR_PATH with TGI
+
+volume=$PWD/data # share a volume with the Docker container to avoid downloading weights every run
+
+docker run -e HUGGING_FACE_HUB_TOKEN=$HUGGINGFACE_AUTH_TOKEN --shm-size 1g -p 8081:80 -v $volume:/data ghcr.io/huggingface/text-generation-inference:1.0.0 --model-id $HUGGINGFACE_MODEL_NAME_OR_PATH