granite

app.py

@@ -2,18 +2,17 @@ from huggingface_hub import InferenceClient
 
 import gradio as gr
 
-client = InferenceClient(
-    "mistralai/Mistral-7B-Instruct-v0.3"
-)
+model_path = "ibm-granite/granite-34b-code-instruct-8k"
+model_path="ibm-granite/granite-3b-code-instruct-2k"
+client = InferenceClient(model_path)
 
-def format_prompt(message, history, system_message=None):
-    prompt = "<s>"
+
+def format_prompt(system, message, history):
+    prompt = [{"role": "system", "content": system}] 
     for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST]"
-        prompt += f" {bot_response}</s> "
-    if system_message:
-        prompt += f"[SYS] {system_message} [/SYS]"
-    prompt += f"[INST] {message} [/INST]"
+        prompt += {"role": "user", "content": user_prompt}
+        prompt += {"role": "assistant", "content": bot_response}
+    prompt += {"role": "user", "content": message}
     return prompt
 
 def generate(
@@ -32,7 +31,7 @@ def generate(
         seed=42,
     )
 
-    formatted_prompt = format_prompt(prompt, history, system_message)
+    formatted_prompt = format_prompt(system_message, prompt, history)
 
     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
     output = ""
@@ -94,5 +93,5 @@ gr.ChatInterface(
     fn=generate,
     chatbot=gr.Chatbot(show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel"),
     additional_inputs=additional_inputs,
-    title="""mistralai/Mistral-7B-Instruct-v0.3"""
+    title=model_path
 ).launch(show_api=True)