feat: implement DeepSeek Janus chat interface

- Add streaming response generation
- Implement chat UI with real-time updates
- Add multi-page structure
- Update dependencies
- Improve error handling

README.md

@@ -1,5 +1,5 @@
 ---
-title: DeepSeek R1 Chat
+title: Janus Pro Chat
 emoji: 🧠
 colorFrom: blue
 colorTo: purple
@@ -7,21 +7,21 @@ sdk: streamlit
 sdk_version: 1.41.1
 app_file: app.py
 pinned: false
-short_description: Advanced Chat using DeepSeek-R1-Distill-8B
+short_description: Advanced Chat using Janus-Pro-7B
 ---
 
-# DeepSeek R1 Chat Assistant
+# Janus Pro Chat Assistant
 
-Powerful chat interface powered by DeepSeek-R1-Distill-Llama-8B model.
+Professional chat interface powered by DeepSeek Janus-Pro-7B model.
 
 ## Features
-- Advanced language understanding
-- Context-aware responses
-- Efficient 8B parameter model
+- Professional-grade responses
+- Real-time streaming
+- Context-aware chat
 - Local CPU inference
 - Memory optimized
 
 ## Usage
-- Type your message and press Enter
-- Clear chat history using sidebar button
-- Best for complex conversations
+- Get detailed, accurate responses
+- Perfect for professional and technical discussions
+- Maintains conversation context

app.py

@@ -5,7 +5,7 @@ import logging
 
 # Configure page
 st.set_page_config(
-    page_title="DeepSeek R1 Chat",
+    page_title="DeepSeek Assistant",
     page_icon="🧠",
     layout="wide",
     initial_sidebar_state="expanded"
@@ -23,9 +23,17 @@ st.markdown("""
 </style>
 """, unsafe_allow_html=True)
 
+st.title("🧠 DeepSeek AI Assistant")
+st.markdown("""
+Welcome! Choose a feature from the sidebar to get started.
+
+- 💭 Chat Interface: Have a conversation with the AI
+- More features coming soon...
+""")
+
 @st.cache_resource
 def load_model():
-    model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
+    model_name = "deepseek-ai/Janus-Pro-7B"
     
     try:
         tokenizer = AutoTokenizer.from_pretrained(
@@ -53,10 +61,10 @@ def load_model():
 
 def generate_response(prompt, model, tokenizer):
     try:
-        chat_prompt = f"""user
-{prompt}
-assistant
-I'll help you with that."""
+        # Janus-Pro specific prompt format
+        chat_prompt = f"""### Human: {prompt}
+
+### Assistant: Let me help you with that."""
         
         inputs = tokenizer(
             chat_prompt,
@@ -72,7 +80,7 @@ I'll help you with that."""
         
         with torch.inference_mode():
             generated_ids = []
-            for i in range(512):  # Max new tokens
+            for _ in range(512):  # Max new tokens
                 # Generate next token
                 outputs = model.generate(
                     inputs["input_ids"] if not generated_ids else torch.cat([inputs["input_ids"], torch.tensor([generated_ids]).to(model.device)], dim=1),
@@ -80,6 +88,7 @@ I'll help you with that."""
                     temperature=0.7,
                     do_sample=True,
                     top_p=0.95,
+                    top_k=50,  # Added for better quality
                     repetition_penalty=1.1,
                     pad_token_id=tokenizer.eos_token_id
                 )
@@ -93,12 +102,13 @@ I'll help you with that."""
                 message_placeholder.markdown(full_response)
                 
                 # Check for end of generation
-                if next_token == tokenizer.eos_token_id:
+                if next_token == tokenizer.eos_token_id or "### Human:" in full_response:
                     break
             
-            # Clean up response
-            response = full_response.split("assistant")[-1].strip()
-            return response.split("user")[0].strip()
+            # Clean up response for Janus format
+            response = full_response.split("### Assistant:")[-1].strip()
+            response = response.split("### Human:")[0].strip()
+            return response
             
     except Exception as e:
         st.error(f"Error: {str(e)}")

pages/01_chat.py

@@ -0,0 +1,79 @@
+import streamlit as st
+from utils import load_model, generate_stream
+import time
+
+def init_chat():
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+        st.session_state.model, st.session_state.tokenizer = load_model()
+        st.session_state.generating = False
+
+# Chat interface with error handling
+try:
+    st.title("💭 Chat Interface")
+    init_chat()
+    
+    # Sidebar controls with session management
+    with st.sidebar:
+        st.markdown("### Chat Controls")
+        cols = st.columns(2)
+        with cols[0]:
+            if st.button("🗑️ Clear Chat", use_container_width=True):
+                st.session_state.messages = []
+                st.rerun()
+        with cols[1]:
+            if st.button("🔄 Reset Model", use_container_width=True):
+                st.cache_resource.clear()
+                st.rerun()
+    
+    # Chat history with proper formatting
+    chat_container = st.container()
+    with chat_container:
+        for msg in st.session_state.messages:
+            with st.chat_message(msg["role"]):
+                st.markdown(msg["content"])
+    
+    # Input handling with safeguards
+    if prompt := st.chat_input(
+        "Ask me anything...",
+        disabled=st.session_state.get("generating", False)
+    ):
+        if not st.session_state.generating:
+            st.session_state.generating = True
+            
+            # Show user message
+            st.session_state.messages.append({"role": "user", "content": prompt})
+            with st.chat_message("user"):
+                st.markdown(prompt)
+            
+            # Generate and show response
+            with st.chat_message("assistant"):
+                try:
+                    # Get recent context
+                    context = "\n".join([
+                        f"{m['role']}: {m['content']}"
+                        for m in st.session_state.messages[-3:]
+                    ])
+                    
+                    response = generate_stream(
+                        context,
+                        st.session_state.model,
+                        st.session_state.tokenizer
+                    )
+                    
+                    if response:
+                        st.session_state.messages.append({
+                            "role": "assistant",
+                            "content": response
+                        })
+                    
+                except Exception as e:
+                    st.error("Failed to generate response. Please try again.")
+                    st.error(f"Error details: {str(e)}")
+                
+                finally:
+                    st.session_state.generating = False
+
+except Exception as e:
+    st.error(f"Application error: {str(e)}")
+    st.button("🔄 Restart App")

requirements.txt

@@ -1,7 +1,33 @@
 # Core dependencies
 streamlit>=1.41.1
+watchdog>=3.0.0
+
+# Machine Learning and Models
 torch>=2.0.0
-transformers>=4.31.0
-accelerate>=0.21.0
+transformers>=4.35.0
+accelerate>=0.25.0
 sentencepiece>=0.1.99
-einops>=0.6.1
+einops>=0.7.0
+bitsandbytes>=0.41.1
+safetensors>=0.4.0
+
+# Performance optimization
+numpy>=1.24.0
+scipy>=1.11.0
+psutil>=5.9.0
+typing-extensions>=4.8.0
+
+# Text processing
+tiktoken>=0.5.1
+regex>=2023.10.3
+tokenizers>=0.15.0
+
+# UI enhancements
+streamlit-chat>=0.1.1
+streamlit-option-menu>=0.3.2
+streamlit-extras>=0.3.4
+markdown>=3.5.1
+
+# Monitoring and logging
+tqdm>=4.66.1
+rich>=13.7.0

utils.py

@@ -0,0 +1,130 @@
+import streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import logging
+from typing import Generator, Optional
+import time
+
+logging.basicConfig(level=logging.INFO)
+
+@st.cache_resource
+def load_model():
+    model_name = "deepseek-ai/Janus-Pro-7B"
+    
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            padding_side='left'
+        )
+        tokenizer.pad_token = tokenizer.eos_token
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float32,
+            low_cpu_mem_usage=True,
+            trust_remote_code=True,
+            device_map='cpu'
+        )
+        
+        model.eval()
+        torch.set_num_threads(8)
+        return model, tokenizer
+        
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        st.stop()
+
+def stream_tokens(response: str, delay: float = 0.01) -> Generator[str, None, None]:
+    """Stream tokens with controlled delay for smooth output"""
+    buffer = ""
+    for char in response:
+        buffer += char
+        if len(buffer) >= 3 or char in '.!?':  # Stream by chunks or punctuation
+            yield buffer
+            buffer = ""
+            time.sleep(delay)
+    if buffer:  # Yield remaining text
+        yield buffer
+
+def generate_stream(prompt: str, model: AutoModelForCausalLM, tokenizer: AutoTokenizer) -> Optional[str]:
+    try:
+        # Format prompt with safety checks
+        safe_prompt = prompt.strip().replace("<", "&lt;").replace(">", "&gt;")
+        chat_prompt = f"""### Human: {safe_prompt}
+
+### Assistant: I'll help you with that."""
+
+        # Create persistent placeholder
+        message_placeholder = st.empty()
+        response_container = st.container()
+        
+        with torch.inference_mode(), st.spinner("Thinking..."):
+            inputs = tokenizer(
+                chat_prompt,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=2048
+            )
+
+            # Stream generation with progress tracking
+            generated_text = ""
+            generated_ids = []
+            progress_bar = st.progress(0)
+            
+            for i in range(512):  # Max tokens
+                try:
+                    outputs = model.generate(
+                        inputs["input_ids"] if not generated_ids else torch.cat([inputs["input_ids"], torch.tensor([generated_ids]).to(model.device)], dim=1),
+                        max_new_tokens=1,
+                        temperature=0.7,
+                        do_sample=True,
+                        top_p=0.95,
+                        top_k=50,
+                        repetition_penalty=1.1,
+                        pad_token_id=tokenizer.eos_token_id,
+                        attention_mask=torch.ones_like(inputs["input_ids"] if not generated_ids else torch.cat([inputs["input_ids"], torch.tensor([generated_ids]).to(model.device)], dim=1))
+                    )
+                    
+                    next_token = outputs[0][-1].item()
+                    generated_ids.append(next_token)
+                    
+                    # Update progress
+                    progress = min(1.0, i / 512)
+                    progress_bar.progress(progress)
+                    
+                    # Decode and stream current output
+                    current_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
+                    
+                    # Stream tokens smoothly
+                    for chunk in stream_tokens(current_text[len(generated_text):]):
+                        generated_text += chunk
+                        with response_container:
+                            message_placeholder.markdown(generated_text)
+                    
+                    # Check stopping conditions
+                    if (next_token == tokenizer.eos_token_id or 
+                        "### Human:" in current_text or 
+                        len(generated_ids) >= 512):
+                        break
+                        
+                except torch.cuda.OutOfMemoryError:
+                    torch.cuda.empty_cache()
+                    st.warning("Memory limit reached, truncating response...")
+                    break
+                    
+            progress_bar.empty()
+            
+            # Clean and validate response
+            response = generated_text.split("### Assistant:")[-1].split("### Human:")[0].strip()
+            if len(response) < 10:  # Minimum response length
+                raise ValueError("Generated response too short")
+                
+            return response
+            
+    except Exception as e:
+        logger = logging.getLogger(__name__)
+        logger.error(f"Generation error: {str(e)}")
+        st.error("Something went wrong. Please try again with a different prompt.")
+        return None