Update app.py

app.py

@@ -1,151 +1,139 @@
-from huggingface_hub import InferenceClient
-import gradio as gr
-import base64
-from PIL import Image
-import io
+import streamlit as st
+import cohere
 
-def image_to_data_url(image_path):
-    if image_path is None:
-        return None
-    with Image.open(image_path) as img:
-        buffered = io.BytesIO()
-        img_format = img.format if img.format else "JPEG"
-        img.save(buffered, format=img_format)
-        img_str = base64.b64encode(buffered.getvalue()).decode()
-        return f"data:image/{img_format.lower()};base64,{img_str}"
+st.set_page_config(page_title="Cohere Chat Interface", layout="wide")
 
-def process_input(image, image_url, prompt, model, hf_token):
-    if not hf_token.startswith("hf_"):
-        raise gr.Error("Invalid Hugging Face token. It should start with 'hf_'")
+MODEL_PFPS = {
+    "command-a-03-2025": "/media/pfp/cohere-pfp.png",
+    "command-r7b-12-2024": "/media/pfp/cohere-pfp.png",
+    "command-r-plus-04-2024": "/media/pfp/cohere-pfp.png",
+    "command-r-plus": "/media/pfp/cohere-pfp.png",
+    "command-r-08-2024": "/media/pfp/cohere-pfp.png",
+    "command-r-03-2024": "/media/pfp/cohere-pfp.png",
+    "command-r": "/media/pfp/cohere-pfp.png",
+    "command": "/media/pfp/cohere-pfp.png",
+    "command-nightly": "/media/pfp/cohere-pfp.png",
+    "command-light": "/media/pfp/cohere-pfp.png",
+    "command-light-nightly": "/media/pfp/cohere-pfp.png"
+}
+
+USER_PFP = "https://example.com/user-default.png"
+
+MODEL_INFO = {
+    "command-a-03-2025": {
+        "description": "Command A is our most performant model to date, excelling at tool use, agents, retrieval augmented generation (RAG), and multilingual use cases. Command A has a context length of 256K, only requires two GPUs to run, and has 150% higher throughput compared to Command R+ 08-2024.",
+        "context_window": "256K tokens",
+        "output_tokens": "8K tokens"
+    },
+    "command-r7b-12-2024": {
+        "description": "command-r7b-12-2024 is a small, fast update delivered in December 2024. It excels at RAG, tool use, agents, and similar tasks requiring complex reasoning and multiple steps.",
+        "context_window": "128K tokens",
+        "output_tokens": "4K tokens"
+    },
+    "command-r-plus-04-2024": {
+        "description": "Command R+ is an instruction-following conversational model that performs language tasks at a higher quality, more reliably, and with a longer context than previous models. It is best suited for complex RAG workflows and multi-step tool use.",
+        "context_window": "128K tokens",
+        "output_tokens": "4K tokens"
+    },
+    "command-r-plus": {
+        "description": "command-r-plus is an alias for command-r-plus-04-2024, so if you use command-r-plus in the API, that's the model you're pointing to.",
+        "context_window": "128K tokens",
+        "output_tokens": "4K tokens"
+    },
+    "command-r-08-2024": {
+        "description": "command-r-08-2024 is an update of the Command R model, delivered in August 2024.",
+        "context_window": "128K tokens",
+        "output_tokens": "4K tokens"
+    },
+    "command-r-03-2024": {
+        "description": "Command R is an instruction-following conversational model that performs language tasks at a higher quality, more reliably, and with a longer context than previous models. It can be used for complex workflows like code generation, retrieval augmented generation (RAG), tool use, and agents.",
+        "context_window": "128K tokens",
+        "output_tokens": "4K tokens"
+    },
+    "command-r": {
+        "description": "command-r is an alias for command-r-03-2024, so if you use command-r in the API, that's the model you're pointing to.",
+        "context_window": "128K tokens",
+        "output_tokens": "4K tokens"
+    },
+    "command": {
+        "description": "An instruction-following conversational model that performs language tasks with high quality, more reliably and with a longer context than our base generative models.",
+        "context_window": "4K tokens",
+        "output_tokens": "4K tokens"
+    },
+    "command-nightly": {
+        "description": "Nightly version of command - experimental and unstable. Not recommended for production use.",
+        "context_window": "128K tokens",
+        "output_tokens": "4K tokens"
+    },
+    "command-light": {
+        "description": "Smaller, faster version of command with similar capabilities.",
+        "context_window": "4K tokens",
+        "output_tokens": "4K tokens"
+    },
+    "command-light-nightly": {
+        "description": "Nightly version of command-light - experimental and unstable. Not for production use.",
+        "context_window": "128K tokens",
+        "output_tokens": "4K tokens"
+    }
+}
+
+with st.sidebar:
+    st.title("Configuration")
+    api_key = st.text_input("Cohere API Key", type="password")
     
-    client = InferenceClient(
-        api_key=hf_token,
-        provider="cohere"
+    selected_model = st.selectbox(
+        "Select Model",
+        options=list(MODEL_INFO.keys()),
+        format_func=lambda x: x.upper()
     )
-
-    image_data = None
-    if image is not None:
-        image_data = image_to_data_url(image)
-    elif image_url:
-        image_data = image_url
     
-    if not image_data:
-        raise gr.Error("Please provide either an image upload or image URL")
+    st.divider()
+    st.subheader("Model Details")
+    st.image(MODEL_PFPS[selected_model], width=80)
+    st.markdown(f"**{selected_model}**")
+    st.markdown(MODEL_INFO[selected_model]["description"])
+    st.markdown(f"**Context Window:** {MODEL_INFO[selected_model]['context_window']}")
+    st.markdown(f"**Max Output:** {MODEL_INFO[selected_model]['output_tokens']}")
 
-    messages = [{
-        "role": "user",
-        "content": [
-            {"type": "text", "text": prompt},
-            {"type": "image_url", "image_url": {"url": image_data}}
-        ]
-    }]
-    
-    try:
-        stream = client.chat.completions.create(
-            model=model,
-            messages=messages,
-            max_tokens=8000,
-            stream=True,
-        )
-        
-        full_response = ""
-        for chunk in stream:
-            if hasattr(chunk.choices[0], 'delta') and hasattr(chunk.choices[0].delta, 'content'):
-                content = chunk.choices[0].delta.content or ""
-                full_response += content
-                yield full_response
-            elif hasattr(chunk, 'content'):
-                content = chunk.content or ""
-                full_response += content
-                yield full_response
-    except Exception as e:
-        raise gr.Error(f"API Error: {str(e)}")
+st.title(f"Chat with {selected_model.upper()}")
+st.image(MODEL_PFPS[selected_model], width=50)
+
+if "messages" not in st.session_state:
+    st.session_state.messages = []
 
-models = [
-    "CohereLabs/aya-vision-32b",
-    "CohereLabs/aya-vision-8b",
-]
+for message in st.session_state.messages:
+    avatar = USER_PFP if message["role"] == "user" else MODEL_PFPS[selected_model]
+    with st.chat_message(message["role"], avatar=avatar):
+        st.markdown(message["content"])
 
-with gr.Blocks() as demo:
-    gr.Markdown("""
-    # 🔍 Aya-Vision Model Interface
+if prompt := st.chat_input("Type your message..."):
+    if not api_key:
+        st.error("API key required - enter in sidebar")
+        st.stop()
     
-    *Explore state-of-the-art vision-language models by Cohere through this interface.  
-    Supports image inputs via upload or URL, with streaming responses.*
-    Read more about Aya Vision [here](https://cohere.com/research/aya)
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user", avatar=USER_PFP):
+        st.markdown(prompt)
     
-    **Get your HF token:** [Hugging Face Settings](https://huggingface.co/settings/tokens)
-    """)
-
-    with gr.Row():
-        with gr.Column():
-            hf_token = gr.Textbox(
-                label="Hugging Face Token",
-                type="password",
-                placeholder="hf_XXXXXXXXXXXXXX",
-                info="Token is used temporarily for the request"
-            )
-            
-            model_choice = gr.Dropdown(
-                label="Model Selection",
-                choices=models,
-                value=models[0]
+    try:
+        co = cohere.ClientV2(api_key)
+        
+        with st.chat_message("assistant", avatar=MODEL_PFPS[selected_model]):
+            response = co.chat(
+                model=selected_model,
+                messages=st.session_state.messages
             )
             
-            with gr.Tab("Upload Image"):
-                image_input = gr.Image(
-                    label="Upload Image",
-                    type="filepath",
-                    sources=["upload"]
-                )
-            with gr.Tab("Image URL"):
-                image_url = gr.Textbox(
-                    label="Image URL",
-                    placeholder="https://example.com/image.jpg",
-                )
+            if hasattr(response, 'text'):
+                full_response = response.text
+            else:
+                full_response = "Error: Unexpected API response format"
             
-            prompt = gr.Textbox(
-                label="Prompt",
-                value="Describe this image in one sentence.",
-                lines=3
-            )
-            submit_btn = gr.Button("Generate", variant="primary")
+            st.markdown(full_response)
         
-        with gr.Column():
-            output = gr.Textbox(
-                label="Model Response",
-                interactive=False,
-                lines=10,
-                autoscroll=True
-            )
-
-    submit_btn.click(
-        fn=process_input,
-        inputs=[image_input, image_url, prompt, model_choice, hf_token],
-        outputs=output,
-        concurrency_limit=None
-    )
-
-    gr.Examples(
-        examples=[
-            [
-                None,
-                "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg",
-                "Describe this image in one sentence.",
-                models[0],
-                ""
-            ],
-            [
-                None,
-                "https://upload.wikimedia.org/wikipedia/commons/4/47/PNG_transparency_demonstration_1.png",
-                "What is unique about this image format?",
-                models[1],
-                ""
-            ]
-        ],
-        inputs=[image_input, image_url, prompt, model_choice, hf_token],
-        label="Try these examples:"
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+        st.session_state.messages.append({"role": "assistant", "content": full_response})
+    
+    except cohere.CohereError as e:
+        st.error(f"Cohere API Error: {str(e)}")
+    except Exception as e:
+        st.error(f"General Error: {str(e)}")