File size: 4,686 Bytes
7743041
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
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()