mlx-community

/

gemma-3n-E4B-it-bf16

Like 12

Follow

MLX Community 17.6k

from mlx_lm import load, generate
import sys
import time  # for optional delay during streaming

model_name = "mlx-community/gemma-3n-E2B-it-lm-bf16"
model, tokenizer = load(model_name)
system_prompt = "You are a helpful assistant. Provide clear, direct answers."

def create_prompt(user_input, conversation_history=None):
    """Create a clean prompt with minimal formatting"""
    if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template is not None:
        try:
            messages = [{"role": "user", "content": user_input}]
            try:
                messages.insert(0, {"role": "system", "content": system_prompt})
                prompt = tokenizer.apply_chat_template(
                    messages,
                    tokenize=False,
                    add_generation_prompt=True
                )
            except:
                messages = [{"role": "user", "content": user_input}]
                prompt = tokenizer.apply_chat_template(
                    messages,
                    tokenize=False,
                    add_generation_prompt=True
                )
            return prompt
        except Exception as e:
            print(f"Chat template failed: {e}")
    return f"User: {user_input}\nAssistant:"

def stream_generate(model, tokenizer, prompt):
    """Generate with streaming simulation - strip prompt first"""
    try:
        response = generate(
            model=model,
            tokenizer=tokenizer,
            prompt=prompt,
            max_tokens=1024,
            verbose=False
        )

        # Strip the prompt from the response first
        if prompt in response:
            clean_response = response.replace(prompt, "").strip()
        else:
            clean_response = response.strip()

        # Now stream the clean response character by character
        for char in clean_response:
            print(char, end="", flush=True)
            # Optional: uncomment for typing effect
            time.sleep(0.01)

        return clean_response
    except Exception as e:
        print(f"[Error] Generation failed: {e}")
        return f"Error: {e}"

print("Chat with the model (type 'exit', 'quit', or 'q' to exit)")
print("=" * 50)

conversation_count = 0
while True:
    try:
        user_input = input(f"\n[{conversation_count + 1}] You: ").strip()
        if user_input.lower() in ['exit', 'quit', 'q', '']:
            print("\nGoodbye!")
            break

        prompt = create_prompt(user_input)
        print(f"\n[{conversation_count + 1}] Assistant: ", end="", flush=True)

        response = stream_generate(model, tokenizer, prompt)

        # Add a newline after streaming is complete
        print()

        conversation_count += 1
    except KeyboardInterrupt:
        print("\n\nChat interrupted. Goodbye!")
        break
    except Exception as e:
        print(f"\nError: {e}")
        continue