app.py

@@ -132,18 +132,6 @@ def bot(chat_history, **kwargs):
         return chat_history
 
 
-def bot_stream(chat_history, **kwargs):
-    try:
-        message = chat_history[-1][0]
-    except Exception as exc:
-        logger.error(f"{chat_history=}: {exc}")
-        raise gr.Error(f"{chat_history=}")
-        # yield chat_history
-    for elm in model.chat_stream(tokenizer, message, chat_history, **kwargs):
-        chat_history[-1] = [message, elm]
-        yield chat_history
-
-
 SYSTEM_PROMPT = "You are a helpful assistant."
 MAX_MAX_NEW_TOKENS = 1024
 MAX_NEW_TOKENS = 128
@@ -158,7 +146,7 @@ class Config:
     top_p: float = 0.9
 
 
-stats_default = SimpleNamespace(llm=None, system_prompt=SYSTEM_PROMPT, config=Config())
+stats_default = SimpleNamespace(llm=model, system_prompt=SYSTEM_PROMPT, config=Config())
 
 theme = gr.themes.Soft(text_size="sm")
 with gr.Blocks(
@@ -167,10 +155,22 @@ with gr.Blocks(
     css=css,
 ) as block:
     stats = gr.State(stats_default)
-
-    def bot_stream_state(chat_history):
-        config = asdict(stats.value.config)
-        return bot_stream(chat_history, **config)
+    if not torch.cuda.is_available():
+        raise gr.Error("GPU not available, cant run. Turn on GPU and restart")
+
+    model_ = stats.value.llm
+    config = stats.value.config
+    model_.generation_config.update(**asdict(config))
+    def bot_stream(chat_history):
+        try:
+            message = chat_history[-1][0]
+        except Exception as exc:
+            logger.error(f"{chat_history=}: {exc}")
+            raise gr.Error(f"{chat_history=}")
+            # yield chat_history
+        for elm in model.chat_stream(tokenizer, message, chat_history):
+            chat_history[-1] = [message, elm]
+            yield chat_history
 
     with gr.Accordion("🎈 Info", open=False):
         gr.Markdown(
@@ -212,7 +212,7 @@ with gr.Blocks(
         queue=True,
         show_progress="full",
         # api_name=None,
-    ).then(bot_stream_state, chatbot, chatbot, queue=True)
+    ).then(bot_stream, chatbot, chatbot, queue=True)
     submit_click_event = submit.click(
         # fn=lambda x, y: ("",) + user(x, y)[1:],  # clear msg
         fn=user,  # clear msg
@@ -221,7 +221,7 @@ with gr.Blocks(
         queue=True,
         show_progress="full",
         # api_name=None,
-    ).then(bot_stream_state, chatbot, chatbot, queue=True)
+    ).then(bot_stream, chatbot, chatbot, queue=True)
     stop.click(
         fn=None,
         inputs=None,