Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+import os
 
 MODELS = {
     "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
@@ -9,8 +10,19 @@ MODELS = {
     "Cohere Command R+": "CohereForAI/c4ai-command-r-plus",
 }
 
+MODELS_REQUIRING_TOKEN = [
+    "meta-llama/Meta-Llama-3.1-8B-Instruct",
+    "CohereForAI/c4ai-command-r-plus"
+]
+
 def get_client(model_name):
-    return InferenceClient(MODELS[model_name])
+    model_id = MODELS[model_name]
+    if model_id in MODELS_REQUIRING_TOKEN:
+        hf_token = os.getenv("HF_TOKEN")
+        if not hf_token:
+            raise ValueError(f"HF_TOKEN environment variable is required for {model_name}")
+        return InferenceClient(model_id, token=hf_token)
+    return InferenceClient(model_id)
 
 def respond(
     message,
@@ -21,7 +33,11 @@ def respond(
     top_p,
     system_message,
 ):
-    client = get_client(model_name)
+    try:
+        client = get_client(model_name)
+    except ValueError as e:
+        return str(e)
+
     messages = [{"role": "system", "content": system_message}]
 
     for human, assistant in chat_history:
@@ -43,6 +59,9 @@ def respond(
         response += token
         yield response
 
+def clear_conversation():
+    return None
+
 with gr.Blocks() as demo:
     gr.Markdown("# Advanced AI Chatbot")
     gr.Markdown("Chat with different language models and customize your experience!")
@@ -66,10 +85,13 @@ with gr.Blocks() as demo:
         with gr.Column(scale=2):
             chatbot = gr.Chatbot()
             msg = gr.Textbox(label="Your message")
-            clear = gr.Button("Clear")
+            with gr.Row():
+                submit_button = gr.Button("Submit")
+                clear_button = gr.Button("Clear")
 
     msg.submit(respond, [msg, chatbot, model_name, max_tokens, temperature, top_p, system_message], chatbot)
-    clear.click(lambda: None, None, chatbot, queue=False)
+    submit_button.click(respond, [msg, chatbot, model_name, max_tokens, temperature, top_p, system_message], chatbot)
+    clear_button.click(clear_conversation, outputs=chatbot, queue=False)
 
 if __name__ == "__main__":
     demo.launch()