Update app.py

app.py

@@ -18,101 +18,88 @@ print("Loading model, please wait...")
 model, tokenizer = load_model()
 print("Model loaded successfully!")
 
-# List of supported languages
-SUPPORTED_LANGUAGES = [
-    "English", "Spanish", "French", "German", "Chinese", 
-    "Japanese", "Russian", "Arabic", "Portuguese", "Italian"
-]
-
-def translate_text(input_text, source_lang, target_lang, max_length=4096):
+def generate_response(message, chat_history, max_length=4096):
     """
-    Translates text from source language to target language using the BitNet model
+    Generates a response from the BitNet model based on the user's message
     """
-    if not input_text.strip():
-        return "Please enter some text to translate."
+    if not message.strip():
+        return "", chat_history
     
-    # Create a translation prompt
-    prompt = f"""Translate the following {source_lang} text to {target_lang}.
+    # Create a chat prompt based on the history and new message
+    full_prompt = ""
+    for user_msg, bot_msg in chat_history:
+        full_prompt += f"User: {user_msg}\nAssistant: {bot_msg}\n\n"
     
-{source_lang} text: {input_text}
-
-{target_lang} translation:"""
+    full_prompt += f"User: {message}\nAssistant:"
 
     # Create inputs for the model
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
     
-    # Generate translation
+    # Generate response
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
             max_new_tokens=max_length,
-            do_sample=False,  # Use greedy decoding for translation
-            temperature=0.1,  # Low temperature for more deterministic output
+            do_sample=True,
+            temperature=0.7,  # Slightly higher temperature for more creative responses
+            top_p=0.95,
         )
     
-    # Extract only the generated part (the translation)
-    translated_text = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+    # Extract only the generated part (the response)
+    response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
     
-    return translated_text.strip()
+    # Update chat history
+    chat_history.append((message, response.strip()))
+    
+    return "", chat_history
 
 # Define the Gradio interface
-def create_translation_interface():
-    with gr.Blocks(title="BitNet Multilingual Translation Tool") as demo:
-        gr.Markdown("# 🌍 BitNet Multilingual Translation Tool")
-        gr.Markdown("A lightweight translation application powered by Microsoft's BitNet b1.58 2B4T model.")
+def create_chat_interface():
+    with gr.Blocks(title="BitNet Chat Assistant") as demo:
+        gr.Markdown("# 💬 BitNet Chat Assistant")
+        gr.Markdown("A lightweight chat application powered by Microsoft's BitNet b1.58 2B4T model.")
         
-        with gr.Row():
-            with gr.Column():
-                source_lang = gr.Dropdown(
-                    choices=SUPPORTED_LANGUAGES, 
-                    value="English", 
-                    label="Source Language"
-                )
-                input_text = gr.Textbox(
-                    lines=5, 
-                    placeholder="Enter text to translate...",
-                    label="Input Text"
-                )
-                
-            with gr.Column():
-                target_lang = gr.Dropdown(
-                    choices=SUPPORTED_LANGUAGES, 
-                    value="Spanish", 
-                    label="Target Language"
-                )
-                output_text = gr.Textbox(
-                    lines=5, 
-                    label="Translated Text"
-                )
+        chatbot = gr.Chatbot(height=400)
+        msg = gr.Textbox(
+            show_label=False,
+            placeholder="Type your message here...",
+            container=False
+        )
         
-        translate_btn = gr.Button("Translate")
-        translate_btn.click(
-            fn=translate_text,
-            inputs=[input_text, source_lang, target_lang],
-            outputs=output_text
+        clear = gr.Button("Clear Conversation")
+        
+        def clear_convo():
+            return "", []
+        
+        msg.submit(
+            fn=generate_response,
+            inputs=[msg, chatbot],
+            outputs=[msg, chatbot]
         )
         
+        clear.click(fn=clear_convo, inputs=[], outputs=[msg, chatbot])
+        
         # Add some example inputs
         examples = [
-            ["Hello, how are you today?", "English", "Spanish"],
-            ["I'd like to learn more about artificial intelligence.", "English", "French"],
-            ["The weather is beautiful today.", "English", "German"],
-            ["Could you please help me find the nearest restaurant?", "English", "Japanese"],
+            ["Hello, how are you today?"],
+            ["Can you tell me about artificial intelligence?"],
+            ["What's your favorite book?"],
+            ["Write a short poem about technology."],
         ]
-        gr.Examples(examples=examples, inputs=[input_text, source_lang, target_lang])
+        gr.Examples(examples=examples, inputs=[msg])
         
         gr.Markdown("""
         ## About
-        This application uses Microsoft's BitNet b1.58 2B4T, a 1-bit Large Language Model, for translation tasks.
+        This application uses Microsoft's BitNet b1.58 2B4T, a 1-bit Large Language Model, for conversational AI.
         The model runs efficiently on consumer hardware due to its 1-bit architecture, offering significant
         advantages in memory usage, energy consumption, and latency.
         
-        Note: Translation quality may vary by language pair. This is a demonstration of the lightweight model's capabilities.
+        Note: This is a demonstration of the lightweight model's capabilities.
         """)
         
     return demo
 
 # Create and launch the Gradio interface
 if __name__ == "__main__":
-    demo = create_translation_interface()
+    demo = create_chat_interface()
     demo.launch(share=True)  # Set share=False if you don't want a public link