Update app.py

app.py

@@ -1,8 +1,11 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from huggingface_hub import InferenceClient
+
+"""
+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("apple/DCLM-Baseline-7B")
 
-tokenizer = AutoTokenizer.from_pretrained("apple/DCLM-Baseline-7B-8k")
-model = AutoModelForCausalLM.from_pretrained("apple/DCLM-Baseline-7B-8k")
 
 def respond(
     message,
@@ -22,22 +25,23 @@ def respond(
 
     messages.append({"role": "user", "content": message})
 
-    prompt = "".join([f"{'[|Human|] ' if msg['role'] == 'user' else '[|AI|] '}{msg['content']}" for msg in messages])
+    response = ""
 
-    inputs = tokenizer(prompt, return_tensors="pt")
-    gen_kwargs = {
-        "max_new_tokens": max_tokens,
-        "top_p": top_p,
-        "temperature": temperature,
-        "do_sample": True,
-        "repetition_penalty": 1.1
-    }
-    with torch.no_grad():
-        output = model.generate(inputs['input_ids'], **gen_kwargs)
-    response = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)[len(prompt):]
+    for message in client.chat_completion(
+        messages,
+        max_tokens=max_tokens,
+        stream=True,
+        temperature=temperature,
+        top_p=top_p,
+    ):
+        token = message.choices[0].delta.content
 
-    yield response
+        response += token
+        yield response
 
+"""
+For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
+"""
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[