Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,249 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import time
+import os
 
-"""
-# Welcome to Streamlit!
+# Page configuration
+st.set_page_config(
+    page_title="Hakim AI Assistant",
+    page_icon="🤖",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
 
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
+# Custom CSS for better UI
+st.markdown("""
+<style>
+    .main-header {
+        text-align: center;
+        color: #2E86AB;
+        font-size: 2.5rem;
+        margin-bottom: 2rem;
+    }
+    .chat-message {
+        padding: 1rem;
+        border-radius: 10px;
+        margin: 1rem 0;
+    }
+    .user-message {
+        background-color: #E3F2FD;
+        border-left: 5px solid #2196F3;
+    }
+    .assistant-message {
+        background-color: #F1F8E9;
+        border-left: 5px solid #4CAF50;
+    }
+    .stTextArea textarea {
+        border-radius: 10px;
+    }
+</style>
+""", unsafe_allow_html=True)
 
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
+@st.cache_resource
+def load_model_and_tokenizer():
+    """Load the model and tokenizer with caching for better performance"""
+    try:
+        with st.spinner("Loading Hakim model... This may take a few minutes on first load."):
+            # Load tokenizer
+            tokenizer = AutoTokenizer.from_pretrained("Rabe3/Hakim")
+            
+            # Load model with appropriate settings
+            model = AutoModelForCausalLM.from_pretrained(
+                "Rabe3/Hakim",
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                device_map="auto" if torch.cuda.is_available() else None,
+                trust_remote_code=True
+            )
+            
+            # Create pipeline
+            text_pipeline = pipeline(
+                "text-generation",
+                model=model,
+                tokenizer=tokenizer,
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                device_map="auto" if torch.cuda.is_available() else None
+            )
+            
+            return tokenizer, model, text_pipeline
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        return None, None, None
 
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
+def generate_response(pipeline, prompt, system_prompt, max_length=512, temperature=0.7, top_p=0.9, do_sample=True):
+    """Generate response using the model pipeline"""
+    try:
+        # Combine system prompt with user input
+        full_prompt = f"{system_prompt}\n\nUser: {prompt}\nAssistant:"
+        
+        # Generate response
+        with st.spinner("Generating response..."):
+            response = pipeline(
+                full_prompt,
+                max_length=max_length,
+                temperature=temperature,
+                top_p=top_p,
+                do_sample=do_sample,
+                pad_token_id=pipeline.tokenizer.eos_token_id,
+                return_full_text=False,
+                num_return_sequences=1
+            )
+        
+        # Extract generated text
+        generated_text = response[0]['generated_text']
+        
+        # Clean up the response (remove the prompt part if it's included)
+        if "Assistant:" in generated_text:
+            generated_text = generated_text.split("Assistant:")[-1].strip()
+        
+        return generated_text
+    
+    except Exception as e:
+        return f"Error generating response: {str(e)}"
 
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
+def main():
+    # Header
+    st.markdown('<h1 class="main-header">🤖 Hakim AI Assistant</h1>', unsafe_allow_html=True)
+    
+    # Load model
+    tokenizer, model, pipeline = load_model_and_tokenizer()
+    
+    if pipeline is None:
+        st.error("Failed to load the model. Please refresh the page and try again.")
+        return
+    
+    # Sidebar for configuration
+    with st.sidebar:
+        st.header("⚙️ Configuration")
+        
+        # System prompt
+        system_prompt = st.text_area(
+            "System Prompt",
+            value="You are Hakim, a helpful AI assistant. You provide accurate, helpful, and informative responses. You communicate clearly and professionally.",
+            height=150,
+            help="This prompt sets the behavior and personality of the AI assistant."
+        )
+        
+        st.divider()
+        
+        # Generation parameters
+        st.subheader("Generation Parameters")
+        
+        max_length = st.slider(
+            "Max Length",
+            min_value=50,
+            max_value=1000,
+            value=512,
+            step=50,
+            help="Maximum length of generated response"
+        )
+        
+        temperature = st.slider(
+            "Temperature",
+            min_value=0.1,
+            max_value=2.0,
+            value=0.7,
+            step=0.1,
+            help="Controls randomness (lower = more focused, higher = more creative)"
+        )
+        
+        top_p = st.slider(
+            "Top P",
+            min_value=0.1,
+            max_value=1.0,
+            value=0.9,
+            step=0.05,
+            help="Controls diversity via nucleus sampling"
+        )
+        
+        do_sample = st.checkbox(
+            "Enable Sampling",
+            value=True,
+            help="Enable sampling for more diverse responses"
+        )
+        
+        st.divider()
+        
+        # Model info
+        st.subheader("ℹ️ Model Information")
+        st.info("**Model:** Rabe3/Hakim\n**Type:** Causal Language Model\n**Framework:** Transformers")
+        
+        # Clear chat button
+        if st.button("🗑️ Clear Chat History", type="secondary"):
+            if 'messages' in st.session_state:
+                st.session_state.messages = []
+            st.rerun()
+    
+    # Initialize chat history
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    
+    # Main chat interface
+    st.header("💬 Chat Interface")
+    
+    # Display chat history
+    for message in st.session_state.messages:
+        if message["role"] == "user":
+            st.markdown(f'<div class="chat-message user-message"><strong>You:</strong> {message["content"]}</div>', unsafe_allow_html=True)
+        else:
+            st.markdown(f'<div class="chat-message assistant-message"><strong>Hakim:</strong> {message["content"]}</div>', unsafe_allow_html=True)
+    
+    # Chat input
+    user_input = st.text_area(
+        "Enter your message:",
+        height=100,
+        placeholder="Type your message here...",
+        key="user_input"
+    )
+    
+    col1, col2 = st.columns([1, 4])
+    
+    with col1:
+        send_button = st.button("📤 Send", type="primary")
+    
+    with col2:
+        if st.button("💡 Example Questions"):
+            examples = [
+                "What is artificial intelligence?",
+                "Can you help me write a short story?",
+                "Explain quantum computing in simple terms",
+                "What are the benefits of renewable energy?"
+            ]
+            st.session_state.user_input = st.selectbox("Choose an example:", [""] + examples)
+    
+    # Process user input
+    if send_button and user_input.strip():
+        # Add user message to history
+        st.session_state.messages.append({"role": "user", "content": user_input})
+        
+        # Generate response
+        response = generate_response(
+            pipeline=pipeline,
+            prompt=user_input,
+            system_prompt=system_prompt,
+            max_length=max_length,
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=do_sample
+        )
+        
+        # Add assistant response to history
+        st.session_state.messages.append({"role": "assistant", "content": response})
+        
+        # Rerun to update the display
+        st.rerun()
+    
+    # Footer
+    st.divider()
+    st.markdown(
+        """
+        <div style='text-align: center; color: #666; margin-top: 2rem;'>
+            <p>Powered by <strong>Rabe3/Hakim</strong> model from Hugging Face 🤗</p>
+            <p><em>This AI assistant is designed to be helpful, harmless, and honest.</em></p>
+        </div>
+        """,
+        unsafe_allow_html=True
+    )
 
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+if __name__ == "__main__":
+    main()