Update app.py

app.py

@@ -12,6 +12,11 @@ model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = model.to(device)
 
+# System content - Define the assistant's personality and capabilities
+SYSTEM_CONTENT = """You are TinyLlama, a friendly and helpful AI assistant. 
+You are based on the TinyLlama-1.1B-Chat model and you excel at providing clear, 
+concise answers to various questions. You are knowledgeable about general topics 
+and always strive to be helpful and accurate in your responses."""
 
 # Defining a custom stopping criteria class for the model's text generation.
 class StopOnTokens(StoppingCriteria):
@@ -24,13 +29,19 @@ class StopOnTokens(StoppingCriteria):
 
 
 # Function to generate model predictions.
-def predict(message, history):
+def predict(message, history, system_content=None):
+    # Use custom system content if provided, otherwise use default
+    current_system_content = system_content if system_content else SYSTEM_CONTENT
+    
     history_transformer_format = history + [[message, ""]]
     stop = StopOnTokens()
 
-    # Formatting the input for the model.
-    messages = "</s>".join(["</s>".join(["\n<|user|>:" + item[0], "\n<|assistant|>:" + item[1]])
+    # Formatting the input for the model with system content
+    system_prompt = f"<|system|>\n{current_system_content}</s>"
+    conversation = "</s>".join(["</s>".join(["\n<|user|>:" + item[0], "\n<|assistant|>:" + item[1]])
                         for item in history_transformer_format])
+    
+    messages = system_prompt + conversation
     model_inputs = tokenizer([messages], return_tensors="pt").to(device)
     streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
@@ -54,9 +65,88 @@ def predict(message, history):
         yield partial_message
 
 
-# Setting up the Gradio chat interface.
-gr.ChatInterface(predict,
-                 title="Tinyllama_chatBot",
-                 description="Ask Tiny llama any questions",
-                 examples=['How to cook a fish?', 'Who is the president of US now?']
-                 ).launch()  # Launching the web interface.
+# Custom function to handle system content updates
+def update_system_content(system_content):
+    global SYSTEM_CONTENT
+    if system_content.strip():
+        SYSTEM_CONTENT = system_content
+        return "System content updated successfully!"
+    else:
+        return "Please enter valid system content."
+
+
+# Additional function to reset system content to default
+def reset_system_content():
+    global SYSTEM_CONTENT
+    default_content = """You are TinyLlama, a friendly and helpful AI assistant. 
+You are based on the TinyLlama-1.1B-Chat model and you excel at providing clear, 
+concise answers to various questions. You are knowledgeable about general topics 
+and always strive to be helpful and accurate in your responses."""
+    SYSTEM_CONTENT = default_content
+    return default_content, "System content reset to default!"
+
+
+# Create the interface with additional components for system content
+with gr.Blocks(title="TinyLlama ChatBot with System Content", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🦙 TinyLlama ChatBot")
+    gr.Markdown("Chat with TinyLlama-1.1B-Chat model. Customize the system content to change how the assistant behaves.")
+    
+    # System content configuration section
+    with gr.Accordion("⚙️ System Content Configuration", open=False):
+        gr.Markdown("Customize the assistant's personality and behavior:")
+        system_content_input = gr.Textbox(
+            label="System Content",
+            value=SYSTEM_CONTENT,
+            lines=4,
+            placeholder="Enter system content that defines the assistant's behavior...",
+            info="This content will shape how the AI assistant responds to your questions."
+        )
+        with gr.Row():
+            update_btn = gr.Button("Update System Content")
+            reset_btn = gr.Button("Reset to Default")
+        system_status = gr.Textbox(label="Status", interactive=False)
+    
+    # Chat interface section
+    with gr.Row():
+        with gr.Column(scale=2):
+            gr.Markdown("### 💬 Chat with TinyLlama")
+            chat_interface = gr.ChatInterface(
+                predict,
+                examples=['How to cook a fish?', 'Who is the president of US now?', 'Explain quantum computing simply'],
+                cache_examples=False
+            )
+        with gr.Column(scale=1):
+            gr.Markdown("### ℹ️ About")
+            gr.Markdown("""
+            **Model:** TinyLlama-1.1B-Chat-v1.0  
+            **Parameters:** 1.1 Billion  
+            **Context Window:** 2048 tokens  
+            
+            **Capabilities:**
+            - General conversation
+            - Question answering
+            - Creative writing
+            - Code generation
+            - And much more!
+            
+            **Tips:**
+            - Use the system content to customize behavior
+            - Be specific in your questions
+            - The model works best with clear, concise prompts
+            """)
+    
+    # Event handlers for system content updates
+    update_btn.click(
+        update_system_content,
+        inputs=[system_content_input],
+        outputs=[system_status]
+    )
+    
+    reset_btn.click(
+        reset_system_content,
+        outputs=[system_content_input, system_status]
+    )
+
+
+if __name__ == "__main__":
+    demo.launch(share=False, server_name="0.0.0.0", server_port=7860)