attempt to use unsloth

app.py

@@ -1,57 +1,100 @@
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
 import gradio as gr
 from huggingface_hub import InferenceClient
+from torch.cuda import is_available
+
+from unsloth import FastLanguageModel
+from transformers import TextIteratorStreamer
+from threading import Thread
 
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
 """
 
-client = InferenceClient()
-
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    model,
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-        model=model,
+# client = InferenceClient()
+class MyModel:
+    def __init__(self):
+        self.client = None
+        self.current_model = ""
+        self.tokenizer = None
+
+    def respond(
+        self,
+        message,
+        history: list[tuple[str, str]],
+        model,
+        system_message,
+        max_tokens,
+        temperature,
+        min_p,
     ):
-        token = message.choices[0].delta.content
+        if model != self.current_model or self.current_model is None:
+            client, tokenizer = FastLanguageModel.from_pretrained(
+                model_name = model,
+                max_seq_length = 2048,
+                dtype = None,
+                load_in_4bit = True,
+            )
+            FastLanguageModel.for_inference(client) # Enable native 2x faster inference
+            self.client = client
+            self.tokenizer = tokenizer
+            self.current_model = model
+        
+        text_streamer = TextIteratorStreamer(self.tokenizer, skip_prompt = True)
+
+        messages = [{"role": "system", "content": system_message}]
+
+        for val in history:
+            if val[0]:
+                messages.append({"role": "user", "content": val[0]})
+            if val[1]:
+                messages.append({"role": "assistant", "content": val[1]})
+
+        messages.append({"role": "user", "content": message})
+
+        inputs = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize = True,
+            add_generation_prompt = True, # Must add for generation
+            return_tensors = "pt",
+        ).to("cuda" if is_available() else "cpu")
+        
+        generation_kwargs = dict(input_ids=inputs, streamer=text_streamer, max_new_tokens=max_tokens, use_cache=True, temperature=temperature, min_p=min_p)
+        thread = Thread(target=self.client.generate, kwargs=generation_kwargs)
+        thread.start()
+
+        response = ""
+
+        for new_text in text_streamer:
+            response += new_text
+            yield response.strip("<|eot_id|>")
+
+        # for message in client.chat_completion(
+        #     messages,
+        #     max_tokens=max_tokens,
+        #     stream=True,
+        #     temperature=temperature,
+        #     top_p=top_p,
+        #     model=model,
+        # ):
+        #     token = message.choices[0].delta.content
 
-        response += token
-        yield response
+        #     response += token
+        #     yield response
 
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
 """
+my_model = MyModel()
 model_choices = [
     "lab2-as/lora_model_gguf",
     "lab2-as/lora_model",
 ]
 demo = gr.ChatInterface(
-    respond,
+    my_model.respond,
     additional_inputs=[
         gr.Dropdown(choices=model_choices, value=model_choices[0], label="Select Model"),
         gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
@@ -62,7 +105,7 @@ demo = gr.ChatInterface(
             maximum=1.0,
             value=0.95,
             step=0.05,
-            label="Top-p (nucleus sampling)",
+            label="Min-p (nucleus sampling)",
         ),
     ],
 )