Update gradio_interface.py

gradio_interface.py

@@ -2,224 +2,132 @@ import gradio as gr
 from ai_model import query_ai_model
 import asyncio
 from config import (
-    SYSTEM_MESSAGE,
-    DEFAULT_MAX_NEW_TOKENS,
+    DEFAULT_MAX_OUTPUT_TOKENS,
     DEFAULT_TEMPERATURE,
     DEFAULT_TOP_P,
     DEFAULT_TOP_K,
 )
 import os
-from typing import List, Tuple, Optional, Union
-from PIL import Image
-
-IMAGE_CACHE_DIRECTORY = "/tmp"
-IMAGE_WIDTH = 512
-CHAT_HISTORY = List[Tuple[Optional[Union[Tuple[str], str]], Optional[str]]]
+from typing import List, Optional
 
 def preprocess_stop_sequences(stop_sequences: str) -> Optional[List[str]]:
     if not stop_sequences:
         return None
     return [sequence.strip() for sequence in stop_sequences.split(",")]
 
-def preprocess_image(image: Image.Image) -> Image.Image:
-    image_height = int(image.height * IMAGE_WIDTH / image.width)
-    return image.resize((IMAGE_WIDTH, image_height))
-
-def cache_pil_image(image: Image.Image) -> str:
-    import uuid
-    image_filename = f"{uuid.uuid4()}.jpeg"
-    os.makedirs(IMAGE_CACHE_DIRECTORY, exist_ok=True)
-    image_path = os.path.join(IMAGE_CACHE_DIRECTORY, image_filename)
-    image.save(image_path, "JPEG")
-    return image_path
-
-def upload(files: Optional[List[gr.File]], chatbot: CHAT_HISTORY) -> CHAT_HISTORY:
-    if files is not None:
-        for file in files:
-            image = Image.open(file.name).convert('RGB')
-            image = preprocess_image(image)
-            image_path = cache_pil_image(image)
-            chatbot.append(((image_path,), None))
-    return chatbot
-
-def user(text_prompt: str, chatbot: CHAT_HISTORY):
-    if text_prompt:
-        chatbot.append((text_prompt, None))
-    return "", chatbot
-
-# Function to handle chatbot conversation
-def gradio_chatbot(
-    google_key: str,
-    files: Optional[List[gr.File]],
-    temperature: float,
-    max_new_tokens: int,
-    stop_sequences: str,
-    top_k: int,
-    top_p: float,
-    chatbot: CHAT_HISTORY
-):
-    if len(chatbot) == 0:
-        return chatbot
+def submit_message(message, chat_history):
+    chat_history = chat_history or []
+    chat_history.append({'role': 'user', 'content': message})
+    return "", chat_history
 
+async def gradio_chatbot(
+    message, chat_history, google_key, temperature, max_output_tokens,
+    stop_sequences, top_k, top_p
+):
+    # message is not used here since it's already added in submit_message
     google_key = google_key if google_key else os.environ.get("GOOGLE_API_KEY")
     if not google_key:
-        raise ValueError(
-            "GOOGLE_API_KEY is not set. Please provide your API key."
-        )
+        return chat_history + [{'role': 'assistant', 'content': 'GOOGLE_API_KEY is not set. Please provide your API key.'}]
 
     genai.configure(api_key=google_key)
+
     stop_sequences_list = preprocess_stop_sequences(stop_sequences)
 
-    # Prepare messages for AI model
-    messages = []
-    for user_message, model_message in chatbot:
-        if isinstance(user_message, tuple):
-            # Handle image inputs if necessary
-            pass
-        elif user_message is not None:
-            messages.append({'role': 'user', 'content': user_message})
-        if model_message is not None:
-            messages.append({'role': 'assistant', 'content': model_message})
-
-    # Handle vision model if files are uploaded
-    use_vision = bool(files)
-    image_files = [file.name for file in files] if files else None
-
-    # Run AI model
-    loop = asyncio.get_event_loop()
-    response_text = loop.run_until_complete(
-        query_ai_model(
-            message=messages[-1]['content'],
-            history=messages[:-1],
-            system_message=SYSTEM_MESSAGE,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            stop_sequences=stop_sequences_list,
-        )
+    # Query the AI model
+    assistant_reply = await query_ai_model(
+        message=chat_history[-1]['content'],  # Last user message
+        history=chat_history[:-1],            # Exclude the last message
+        max_output_tokens=max_output_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        stop_sequences=stop_sequences_list,
     )
-    chatbot[-1] = (chatbot[-1][0], response_text)
-    return chatbot
 
-def create_gradio_interface():
-    google_key_component = gr.Textbox(
-        label="GOOGLE API KEY",
-        value="",
-        type="password",
-        placeholder="...",
-        info="You have to provide your own GOOGLE_API_KEY for this app to function properly",
-        visible=os.environ.get("GOOGLE_API_KEY") is None
-    )
-    chatbot_component = gr.Chatbot(
-        label='Gemini',
-        avatar_images=None,
-        height=400,
-        type='messages'  # Use the 'messages' type for better compatibility
-    )
-    text_prompt_component = gr.Textbox(
-        placeholder="Hi there! [press Enter]", show_label=False
-    )
-    upload_button_component = gr.UploadButton(
-        label="Upload Images", file_count="multiple", file_types=["image"]
-    )
-    temperature_component = gr.Slider(
-        minimum=0,
-        maximum=1.0,
-        value=DEFAULT_TEMPERATURE,
-        step=0.05,
-        label="Temperature",
-        info="Controls the randomness of the AI's response."
-    )
-    max_output_tokens_component = gr.Slider(
-        minimum=1,
-        maximum=2048,
-        value=DEFAULT_MAX_NEW_TOKENS,
-        step=1,
-        label="Max new tokens",
-        info="Controls the length of the AI's response."
-    )
-    stop_sequences_component = gr.Textbox(
-        label="Add stop sequence",
-        value="",
-        type="text",
-        placeholder="STOP, END",
-        info="Stops the AI response when any of these sequences are generated."
-    )
-    top_k_component = gr.Slider(
-        minimum=1,
-        maximum=40,
-        value=DEFAULT_TOP_K,
-        step=1,
-        label="Top-K",
-        info="Limits the next token selection to the K most probable tokens."
-    )
-    top_p_component = gr.Slider(
-        minimum=0,
-        maximum=1,
-        value=DEFAULT_TOP_P,
-        step=0.01,
-        label="Top-P",
-        info="Limits the next token selection to tokens with cumulative probability up to P."
-    )
+    # Append assistant's reply to the history
+    chat_history.append(assistant_reply)
 
-    user_inputs = [
-        text_prompt_component,
-        chatbot_component
-    ]
-
-    bot_inputs = [
-        google_key_component,
-        upload_button_component,
-        temperature_component,
-        max_output_tokens_component,
-        stop_sequences_component,
-        top_k_component,
-        top_p_component,
-        chatbot_component
-    ]
+    return chat_history
 
+def create_gradio_interface():
     with gr.Blocks() as demo:
-        gr.HTML("<h1 align='center'>Gemini Playground 💬</h1>")
-        gr.HTML("<h2 align='center'>Play with Gemini Pro and Gemini Pro Vision API</h2>")
-        gr.HTML("""
-        <div style="text-align: center; display: flex; justify-content: center; align-items: center;">
-            <a href="https://huggingface.co/spaces/SkalskiP/ChatGemini?duplicate=true">
-                <img src="https://bit.ly/3gLdBN6" alt="Duplicate Space" style="margin-right: 10px;">
-            </a>
-            <span>Duplicate the Space and run securely with your 
-                <a href="https://makersuite.google.com/app/apikey">GOOGLE API KEY</a>.
-            </span>
-        </div>
-        """)
-        with gr.Column():
-            google_key_component.render()
-            chatbot_component.render()
-            with gr.Row():
-                text_prompt_component.render()
-                upload_button_component.render()
-            with gr.Accordion("Parameters", open=False):
-                temperature_component.render()
-                max_output_tokens_component.render()
-                stop_sequences_component.render()
-                with gr.Accordion("Advanced", open=False):
-                    top_k_component.render()
-                    top_p_component.render()
-
-        text_prompt_component.submit(
-            fn=user,
-            inputs=user_inputs,
-            outputs=[text_prompt_component, chatbot_component],
-            queue=False
-        ).then(
-            fn=gradio_chatbot, inputs=bot_inputs, outputs=[chatbot_component],
+        gr.HTML("<h1 align='center'>Gemini Chat Interface 💬</h1>")
+        gr.HTML("<h2 align='center'>Interact with the Gemini AI model.</h2>")
+
+        google_key_component = gr.Textbox(
+            label="GOOGLE API KEY",
+            value="",
+            type="password",
+            placeholder="Enter your Google API key",
+            visible=os.environ.get("GOOGLE_API_KEY") is None
         )
 
-        upload_button_component.upload(
-            fn=upload,
-            inputs=[upload_button_component, chatbot_component],
-            outputs=[chatbot_component],
-            queue=False
+        chatbot_component = gr.Chatbot(
+            label='Gemini Chatbot',
+            height=600,
+            type='messages'
+        )
+
+        with gr.Row():
+            msg = gr.Textbox(
+                label="Your Message",
+                placeholder="Type your message here and press Enter"
+            )
+
+        with gr.Accordion("Advanced Parameters", open=False):
+            temperature_component = gr.Slider(
+                minimum=0,
+                maximum=1.0,
+                value=DEFAULT_TEMPERATURE,
+                step=0.05,
+                label="Temperature"
+            )
+            max_output_tokens_component = gr.Slider(
+                minimum=1,
+                maximum=2048,
+                value=DEFAULT_MAX_OUTPUT_TOKENS,
+                step=1,
+                label="Max Output Tokens"
+            )
+            stop_sequences_component = gr.Textbox(
+                label="Stop Sequences",
+                value="",
+                placeholder="Comma-separated stop sequences"
+            )
+            top_k_component = gr.Slider(
+                minimum=1,
+                maximum=40,
+                value=DEFAULT_TOP_K,
+                step=1,
+                label="Top-K"
+            )
+            top_p_component = gr.Slider(
+                minimum=0,
+                maximum=1,
+                value=DEFAULT_TOP_P,
+                step=0.01,
+                label="Top-P"
+            )
+
+        state = gr.State()
+
+        msg.submit(
+            fn=submit_message,
+            inputs=[msg, chatbot_component],
+            outputs=[msg, chatbot_component]
+        ).then(
+            fn=gradio_chatbot,
+            inputs=[
+                msg,
+                chatbot_component,
+                google_key_component,
+                temperature_component,
+                max_output_tokens_component,
+                stop_sequences_component,
+                top_k_component,
+                top_p_component
+            ],
+            outputs=chatbot_component
         )
 
     return demo