Update space

app.py

@@ -1,71 +1,154 @@
 import torch
-import gradio as gr
 from transformers import pipeline
-import logging
-import warnings
-from threading import Lock
-
-# Suppress non-critical warnings
-warnings.filterwarnings("ignore", message=".*Torch was not compiled with flash attention.*")
-
-# Setup logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Model configuration
-MODEL_ID = "ezcz/bright-llama-3b-chat"
-MAX_NEW_TOKENS = 256
-TEMPERATURE = 0.1
-TOP_P = 0.9
-TOP_K = 60
-REPETITION_PENALTY = 1.0
-
-# Check for GPU availability
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-logger.info(f"Using device: {device}")
-
-# Load the pipeline
-logger.info("Loading model pipeline...")
-pipe = pipeline(
-    "text-generation",
-    model=MODEL_ID,
-    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto",
-)
+import gradio as gr
+from datetime import datetime
 
-# Define the chat interface
-def chat_interface(user_input, history=None):
-    if history is None:
-        history = []
-    
-    messages = [{"role": "system", "content": ""}, {"role": "user", "content": user_input}]
-    
-    try:
-        outputs = pipe(
-            messages,
-            max_new_tokens=MAX_NEW_TOKENS,
-            temperature=TEMPERATURE,
-            top_p=TOP_P,
-            top_k=TOP_K,
-            repetition_penalty=REPETITION_PENALTY,
+class BrightLlamaChatbot:
+    def __init__(self):
+        self.model_id = "ezcz/bright-llama-3b-chat"
+        self.pipe = pipeline(
+            "text-generation",
+            model=self.model_id,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
         )
-        response = outputs[0]["generated_text"]
-        history.append((user_input, response))
-        return "", history
-    except Exception as e:
-        logger.error(f"Error during response generation: {e}")
-        return "Error generating response.", history
-
-# Define the Gradio interface
-with gr.Blocks() as demo:
-    chatbot = gr.Chatbot()
-    user_input = gr.Textbox(placeholder="Type your message...")
-    clear_button = gr.Button("Clear Chat")
-    submit_button = gr.Button("Send")
+        self.system_prompt = """You are a helpful AI assistant focused on coding and reasoning tasks. 
+        You provide clear, accurate responses while maintaining a friendly tone."""
+        
+    def format_message(self, role, content):
+        timestamp = datetime.now().strftime("%H:%M:%S")
+        return f"<div class='message {role}'><span class='time'>[{timestamp}]</span>{content}</div>"
+    
+    def generate_response(self, user_input, chat_history):
+        if not user_input:
+            return chat_history
+            
+        # Format conversation history for the model
+        messages = [{"role": "system", "content": self.system_prompt}]
+        for msg in chat_history.split("\n"):
+            if msg.strip():
+                if "User:" in msg:
+                    messages.append({"role": "user", "content": msg.replace("User:", "").strip()})
+                elif "Assistant:" in msg:
+                    messages.append({"role": "assistant", "content": msg.replace("Assistant:", "").strip()})
+        
+        messages.append({"role": "user", "content": user_input})
+        
+        # Generate response
+        response = self.pipe(messages, max_new_tokens=512, return_full_text=False, temperature=0.7, top_p=0.9)
+        assistant_response = response[0]["generated_text"]
+        
+        # Format and update chat history
+        updated_history = chat_history + "\n" + f"User: {user_input}" + "\n" + f"Assistant: {assistant_response}"
+        
+        return updated_history
 
-    submit_button.click(chat_interface, inputs=[user_input, chatbot], outputs=[user_input, chatbot])
-    user_input.submit(chat_interface, inputs=[user_input, chatbot], outputs=[user_input, chatbot])
-    clear_button.click(lambda: ([], ""), inputs=[], outputs=[chatbot, user_input])
+    def create_interface(self):
+        with gr.Blocks(css=self.get_custom_css()) as interface:
+            gr.HTML("<h1>🦙 Bright Llama Chatbot</h1>")
+            gr.HTML("<p>An AI assistant specialized in coding and reasoning tasks.</p>")
+            
+            with gr.Row():
+                with gr.Column(scale=4):
+                    chatbot = gr.Textbox(
+                        show_label=False,
+                        placeholder="Conversation history...",
+                        lines=15,
+                        max_lines=15,
+                        interactive=False
+                    )
+                    
+                    with gr.Row():
+                        with gr.Column(scale=8):
+                            user_input = gr.Textbox(
+                                show_label=False,
+                                placeholder="Type your message here...",
+                                lines=2
+                            )
+                        with gr.Column(scale=1):
+                            submit_btn = gr.Button("Send", variant="primary")
+                            clear_btn = gr.Button("Clear")
+                
+                with gr.Column(scale=1):
+                    gr.HTML("<h3>Features</h3>")
+                    gr.HTML("""
+                        <ul>
+                            <li>💬 Natural conversation</li>
+                            <li>🎨 Code syntax highlighting</li>
+                            <li>⚡ Quick responses</li>
+                        </ul>
+                    """)
+                    
+                    with gr.Accordion("Settings", open=False):
+                        temperature = gr.Slider(
+                            minimum=0.1,
+                            maximum=1.0,
+                            value=0.7,
+                            step=0.1,
+                            label="Temperature"
+                        )
+                        
+                        max_length = gr.Slider(
+                            minimum=64,
+                            maximum=1024,
+                            value=512,
+                            step=64,
+                            label="Max Length"
+                        )
+            
+            submit_btn.click(
+                fn=self.generate_response,
+                inputs=[user_input, chatbot],
+                outputs=chatbot
+            ).then(
+                fn=lambda: "",
+                outputs=user_input
+            )
+            
+            clear_btn.click(
+                fn=lambda: ("", ""),
+                outputs=[user_input, chatbot]
+            )
+            
+        return interface
+    
+    def get_custom_css(self):
+        return """
+            .message {
+                padding: 10px;
+                margin: 5px;
+                border-radius: 8px;
+            }
+            .user {
+                background-color: #f0f0f0;
+                margin-left: 20px;
+            }
+            .assistant {
+                background-color: #e3f2fd;
+                margin-right: 20px;
+            }
+            .time {
+                color: #666;
+                font-size: 0.8em;
+                margin-right: 10px;
+            }
+            pre {
+                background-color: #2b2b2b;
+                color: #ffffff;
+                padding: 10px;
+                border-radius: 5px;
+                overflow-x: auto;
+            }
+            code {
+                font-family: 'Courier New', monospace;
+            }
+        """
 
-# Launch the UI
-demo.queue().launch(debug=True, share=True)
+if __name__ == "__main__":
+    chatbot = BrightLlamaChatbot()
+    interface = chatbot.create_interface()
+    interface.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True
+    )