Update app.py

app.py

@@ -0,0 +1,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()