Update app.py

app.py

@@ -1,67 +1,163 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
-
+from openai import OpenAI
 from prompt_template import PromptTemplate, PromptLoader
 from assistant import AIAssistant
+from pathlib import Path
+
+# Load prompts from YAML
+prompts = PromptLoader.load_prompts("prompts.yaml")
 
-"""
-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("HuggingFaceH4/zephyr-7b-beta")
+# Available models and their configurations
+MODELS = {
+    "Zephyr 7B Beta": {
+        "name": "HuggingFaceH4/zephyr-7b-beta",
+        "provider": "huggingface"
+    },
+    "Mistral 7B": {
+        "name": "mistralai/Mistral-7B-v0.1",
+        "provider": "huggingface"
+    },
+    "GPT-3.5 Turbo": {
+        "name": "gpt-3.5-turbo",
+        "provider": "openai"
+    }
+}
 
+# Available prompt strategies
+PROMPT_STRATEGIES = {
+    "Default": "system_context",
+    "Chain of Thought": "cot_prompt",
+    "Knowledge-based": "knowledge_prompt",
+    "Few-shot Learning": "few_shot_prompt",
+    "Meta-prompting": "meta_prompt"
+}
+
+def create_assistant(model_name):
+    model_info = MODELS[model_name]
+    if model_info["provider"] == "huggingface":
+        client = InferenceClient(model_info["name"])
+    else:  # OpenAI
+        client = OpenAI()
+    
+    return AIAssistant(
+        client=client,
+        model=model_info["name"]
+    )
 
 def respond(
     message,
     history: list[tuple[str, str]],
+    model_name,
+    prompt_strategy,
     system_message,
+    override_params: bool,
     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]})
-
+    assistant = create_assistant(model_name)
+    
+    # Get prompt template
+    prompt_template: PromptTemplate = prompts[PROMPT_STRATEGIES[prompt_strategy]]
+    
+    # Generate system message using prompt template
+    formatted_system_message = prompt_template.format(prompt_strategy=system_message)
+    
+    # Prepare messages
+    messages = [{"role": "system", "content": formatted_system_message}]
+    for user_msg, assistant_msg in history:
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if assistant_msg:
+            messages.append({"role": "assistant", "content": assistant_msg})
     messages.append({"role": "user", "content": message})
 
-    response = ""
+    # Get generation parameters
+    generation_params = prompt_template.parameters if not override_params else {
+        "max_tokens": max_tokens,
+        "temperature": temperature,
+        "top_p": top_p
+    }
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
+    # Generate response using the assistant
+    for response in assistant.generate_response(
+        prompt_template=prompt_template,
+        generation_params=generation_params,
         stream=True,
-        temperature=temperature,
-        top_p=top_p,
+        messages=messages
     ):
-        token = message.choices[0].delta.content
-
-        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=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
-
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            model_dropdown = gr.Dropdown(
+                choices=list(MODELS.keys()),
+                value=list(MODELS.keys())[0],
+                label="Select Model"
+            )
+            prompt_strategy_dropdown = gr.Dropdown(
+                choices=list(PROMPT_STRATEGIES.keys()),
+                value=list(PROMPT_STRATEGIES.keys())[0],
+                label="Select Prompt Strategy"
+            )
+            system_message = gr.Textbox(
+                value="You are a friendly and helpful AI assistant.",
+                label="System Message"
+            )
+            
+    with gr.Row():
+        override_params = gr.Checkbox(
+            label="Override Template Parameters",
+            value=False
+        )
+    
+    with gr.Row():
+        with gr.Column(visible=False) as param_controls:
+            max_tokens = gr.Slider(
+                minimum=1,
+                maximum=2048,
+                value=512,
+                step=1,
+                label="Max new tokens"
+            )
+            temperature = gr.Slider(
+                minimum=0.1,
+                maximum=4.0,
+                value=0.7,
+                step=0.1,
+                label="Temperature"
+            )
+            top_p = gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                value=0.95,
+                step=0.05,
+                label="Top-p (nucleus sampling)"
+            )
+    
+    chatbot = gr.ChatInterface(
+        fn=respond,
+        additional_inputs=[
+            model_dropdown,
+            prompt_strategy_dropdown,
+            system_message,
+            override_params,
+            max_tokens,
+            temperature,
+            top_p,
+        ]
+    )
+    
+    def toggle_param_controls(override):
+        return gr.Column(visible=override)
+    
+    override_params.change(
+        toggle_param_controls,
+        inputs=[override_params],
+        outputs=[param_controls]
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()