fixed bugs for tgi

.env.example

@@ -30,7 +30,7 @@ DISABLE_MODEL_PRELOADING=true
 CHAT_HISTORY_ENABLED=true
 SHOW_PARAM_SETTINGS=false
 SHARE_GRADIO_APP=false
-PDF_FILE_BASE_URL=https://ai-engd.netlify.app/pdfs/books/
+PDF_FILE_BASE_URL=https://chat-with-llama-2.netlify.app/pdfs/books/
 
 # if unset, default to "hkunlp/instructor-xl"
 HF_EMBEDDINGS_MODEL_NAME="hkunlp/instructor-large"

app_modules/llm_inference.py

@@ -51,7 +51,6 @@ class LLMInference(metaclass=abc.ABCMeta):
                     streaming_handler,
                 )
                 if streaming_handler is not None
-                and self.llm_loader.streamer.for_huggingface
                 else chain(inputs)
             )
 
@@ -82,20 +81,23 @@ class LLMInference(metaclass=abc.ABCMeta):
         )
         t.start()
 
-        count = (
-            2 if "chat_history" in inputs and len(inputs.get("chat_history")) > 0 else 1
-        )
+        if self.llm_loader.streamer.for_huggingface:
+            count = (
+                2
+                if "chat_history" in inputs and len(inputs.get("chat_history")) > 0
+                else 1
+            )
 
-        while count > 0:
-            try:
-                for token in self.llm_loader.streamer:
-                    streaming_handler.on_llm_new_token(token)
+            while count > 0:
+                try:
+                    for token in self.llm_loader.streamer:
+                        streaming_handler.on_llm_new_token(token)
 
-                self.llm_loader.streamer.reset()
-                count -= 1
-            except Exception:
-                print("nothing generated yet - retry in 0.5s")
-                time.sleep(0.5)
+                    self.llm_loader.streamer.reset()
+                    count -= 1
+                except Exception:
+                    print("nothing generated yet - retry in 0.5s")
+                    time.sleep(0.5)
 
         t.join()
         return que.get()

app_modules/llm_loader.py

@@ -66,9 +66,6 @@ class TextIteratorStreamer(TextStreamer, StreamingStdOutCallbackHandler):
         self.text_queue.put("\n", timeout=self.timeout)
         self.text_queue.put(self.stop_signal, timeout=self.timeout)
 
-    def for_huggingface(self) -> bool:
-        return self.tokenizer != ""
-
     def __iter__(self):
         return self