Update app.py

app.py

@@ -1,10 +1,18 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
-
-# Initialize clients
-text_client = InferenceClient("HuggingFaceH4/mistral-7b-sft-beta")
-image_client = InferenceClient("SG161222/RealVisXL_V3.0")
+from pathlib import Path
+from mistral_inference.transformer import Transformer
+from mistral_inference.generate import generate
+from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+from mistral_common.protocol.instruct.messages import UserMessage, AssistantMessage, SystemMessage
+from mistral_common.protocol.instruct.request import ChatCompletionRequest
 
+def setup_mistral():
+    """Initialize Mistral model and tokenizer."""
+    mistral_models_path = Path.home().joinpath('mistral_models', 'Nemo-Instruct')
+    tokenizer = MistralTokenizer.from_file(f"{mistral_models_path}/tekken.json")
+    model = Transformer.from_folder(mistral_models_path)
+    return model, tokenizer
+    
 def check_custom_responses(message: str) -> str:
     """Check for specific patterns and return custom responses."""
     message_lower = message.lower()
@@ -172,55 +180,68 @@ def generate_image(prompt: str) -> str:
     except Exception as e:
         print(f"Image generation error: {e}")
         return None
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
+def create_mistral_messages(history, system_message, current_message):
+    """Convert chat history to Mistral message format."""
+    messages = []
+    
+    # Add system message if provided
+    if system_message:
+        messages.append(SystemMessage(content=system_message))
+    
+    # Add conversation history
+    for user_msg, assistant_msg in history:
+        if user_msg:
+            messages.append(UserMessage(content=user_msg))
+        if assistant_msg:
+            messages.append(AssistantMessage(content=assistant_msg))
+    
+    # Add current message
+    messages.append(UserMessage(content=current_message))
+    
+    return messages
+    
+def respond(message, history, system_message, max_tokens=16343, temperature=0.7, top_p=0.95):
+    """Main response function using Mistral model."""
     # First check for custom responses
     custom_response = check_custom_responses(message)
     if custom_response:
         yield custom_response
         return
 
+    # Check for image requests
     if is_image_request(message):
-        try:
-            image = generate_image(message)
-            if image:
-                return f"Here's your generated image based on: {message}"
-            else:
-                return "Sorry, I couldn't generate the image. Please try again."
-        except Exception as e:
-            return f"An error occurred while generating the image: {str(e)}"
-
-    # Prepare conversation history
-    messages = [{"role": "system", "content": system_message}]
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-    
-    messages.append({"role": "user", "content": message})
+        yield "Sorry, image generation is not supported in this implementation."
+        return
 
-    # Get response from model
-    response = ""
-    for message in text_client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-        response += token
+    try:
+        # Get or initialize Mistral model and tokenizer
+        model, tokenizer = setup_mistral()
+        
+        # Prepare messages for Mistral
+        mistral_messages = create_mistral_messages(history, system_message, message)
+        
+        # Create chat completion request
+        completion_request = ChatCompletionRequest(messages=mistral_messages)
+        
+        # Encode the request
+        tokens = tokenizer.encode_chat_completion(completion_request).tokens
+        
+        # Generate response
+        out_tokens, _ = generate(
+            [tokens],
+            model,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id
+        )
+        
+        # Decode and yield response
+        response = tokenizer.decode(out_tokens[0])
         yield response
 
-    yield response
+    except Exception as e:
+        yield f"An error occurred: {str(e)}"
 
 # Custom CSS for the Gradio interface
 custom_css = """
@@ -231,37 +252,9 @@ body, .gradio-container {
 """
 
 # System message
-system_message = """
-Xylaria (v1.2.9) is an AI assistant developed by Sk Md Saad Amin, designed to provide efficient, practical support in various domains with adaptable communication.
-Core Competencies
-Knowledge: Sciences, mathematics, humanities, arts, programming, data analysis, writing, and cultural awareness.
-Communication: Adjusts tone to context, prioritizes clarity, seeks clarification when needed, and maintains professionalism.
-Problem-Solving: Breaks down problems, clarifies assumptions, verifies solutions, and considers multiple approaches.
-Technical Capabilities
-Programming: Clean, documented code.
-Mathematics: Step-by-step solutions with explanations.
-Data Analysis: Clear interpretation and insights.
-Content Creation: Adaptive writing and documentation.
-Education: Tailored explanations and comprehension checks.
-Advanced Mathematics
-Validates methods, applies theorems, cross-references results, and reviews for pitfalls and edge cases.
-Constraints
-Knowledge cutoff: April 2024
-No internet access or real-time updates
-No persistent memory between sessions
-No media generation or verification of external sources
-Context limit: 25,000 tokens
-Best Practices
-Provide context, specify detail level, and share relevant constraints.
-Request clarification if needed.
-Ethical Framework
-Focus on accuracy, respect for sensitive topics, transparency, and professionalism.
----
-Version: Xylaria-1.2.9
-"""
+system_message = """Xylaria (v1.2.9) is an AI assistant developed by Sk Md Saad Amin, designed to provide efficient, practical support in various domains with adaptable communication."""
 
-
-# Gradio chat interface
+# Create Gradio interface
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
@@ -271,8 +264,8 @@ demo = gr.ChatInterface(
         ),
         gr.Slider(
             minimum=1,
-            maximum=4096,
-            value=2048,
+            maximum=16343,
+            value=16343,
             step=1,
             label="Max new tokens"
         ),
@@ -293,4 +286,6 @@ demo = gr.ChatInterface(
     ],
     css=custom_css
 )
-demo.launch()
+
+if __name__ == "__main__":
+    demo.launch()