Update app.py

app.py

@@ -1,6 +1,18 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
+import torch
 
+# Load your model once at startup
+print("Loading model...")
+tokenizer = AutoTokenizer.from_pretrained("drumwell/autotrain-2duhi-5mmyz")
+model = AutoModelForCausalLM.from_pretrained(
+    "drumwell/autotrain-2duhi-5mmyz",
+    device_map="auto",
+    torch_dtype=torch.float16,
+    load_in_8bit=True,
+)
+print("Model loaded!")
 
 def respond(
     message,
@@ -14,32 +26,35 @@ def respond(
     """
     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(token=hf_token.token, model="drumwell/autotrain-2duhi-5mmyz")
-
     messages = [{"role": "system", "content": system_message}]
-
     messages.extend(history)
-
     messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
+    
+    # Apply chat template
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(text, return_tensors="pt").to(model.device)
+    
+    # Setup streaming
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = dict(
+        inputs,
+        streamer=streamer,
+        max_new_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
-    ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
-
+        do_sample=True,
+        repetition_penalty=1.1,
+    )
+    
+    # Generate in a separate thread for streaming
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    
+    response = ""
+    for token in streamer:
         response += token
         yield response
 
-
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
 """
@@ -65,6 +80,5 @@ with gr.Blocks() as demo:
         gr.LoginButton()
     chatbot.render()
 
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()