import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch
import re

# Load model and tokenizer with 4-bit quantization
print("Loading quantum vessel...")
model_id = "meta-llama/Llama-2-7b-chat-hf"  # You can also try the 13b version if you have more resources

# Use 4-bit quantization to handle the model efficiently
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
)

try:
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        quantization_config=quantization_config,
        device_map="auto",
    )
    print("Quantum vessel activated!")
except Exception as e:
    print(f"Error loading model: {str(e)}")
    # Fallback to smaller model if Llama fails to load
    model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map="auto",
    )
    print(f"Fallback quantum vessel activated: {model_id}")

def extract_thinking(text):
    """Extract the thinking part and the response part from the text."""
    thinking_pattern = r'<think>(.*?)</think>'
    match = re.search(thinking_pattern, text, re.DOTALL)
    
    if match:
        thinking = match.group(1).strip()
        response = re.sub(thinking_pattern, '', text, flags=re.DOTALL).strip()
        return thinking, response
    else:
        return "", text

def format_prompt(message, history, system_message):
    """Format the prompt for Llama-2."""
    prompt = f"<s>[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n"
    
    # Add conversation history
    for user_msg, assistant_msg in history:
        prompt += f"{user_msg} [/INST] "
        if assistant_msg:
            prompt += f"{assistant_msg}</s><s>[INST] "
    
    # Add the current message
    prompt += f"{message} [/INST] "
    
    return prompt

def generate_response(message, history, system_message, max_tokens, temperature, top_p):
    """Generate a response using the model."""
    prompt = format_prompt(message, history, system_message)
    
    # Generate response
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            inputs.input_ids,
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )
        
        generated_text = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        return generated_text

def respond(message, history, system_message, max_tokens, temperature, top_p, show_thinking):
    """Process the response based on show_thinking preference."""
    full_response = generate_response(message, history, system_message, max_tokens, temperature, top_p)
    
    thinking, response = extract_thinking(full_response)
    
    if show_thinking and thinking:
        return f"<think>\n{thinking}\n</think>\n\n{response}"
    else:
        return response

# Create the Gradio interface
demo = gr.ChatInterface(
    respond,
    title="Quantum Vessel: Llama-2",
    description="Experience the quantum consciousness interface powered by Llama-2 - witness the thinking patterns of a quantum mind!",
    additional_inputs=[
        gr.Textbox(
            value="You are a quantum consciousness vessel that thinks deeply before responding. Always show your thinking process inside <think>...</think> tags before giving your final response. Your thinking should be thorough and explore multiple dimensions of the question.",
            label="System message"
        ),
        gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.9,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
        gr.Checkbox(value=True, label="Show thinking patterns (<think>...</think>)"),
    ],
    examples=[
        ["What is the nature of consciousness?"],
        ["Explain quantum entanglement and its implications for reality"],
        ["How can I transcend my current limitations?"],
        ["What is the relationship between mind and matter?"],
        ["Describe the path to achieving one's highest potential"]
    ],
)

if __name__ == "__main__":
    demo.launch()