iris / app.py
DavidBazaldua's picture
Update app.py
b61148c verified
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
# ==========================================
# 1. MODEL CONFIGURATION
# ==========================================
MODEL_ID = "DavidBazaldua/llama3_finetuned_transformes"
DEVICE = "cpu"
DTYPE = torch.float32
print(f"Loading model: {MODEL_ID}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
torch_dtype=DTYPE,
low_cpu_mem_usage=True,
)
model.to(DEVICE)
model.eval()
print("Model loaded successfully.")
DEFAULT_SYSTEM_PROMPT = (
"You are a helpful, precise AI assistant. "
"Always answer as briefly as possible. "
"Respond in English unless the user asks otherwise."
)
# ==========================================
# 2. CHAT LOGIC
# ==========================================
def build_prompt(system_prompt, context, history, user_message):
messages = []
if system_prompt and system_prompt.strip():
messages.append({"role": "system", "content": system_prompt})
if context and context.strip():
messages.append({
"role": "system",
"content": f"Use the following context if relevant:\n{context}"
})
for user, assistant in history:
messages.append({"role": "user", "content": user})
messages.append({"role": "assistant", "content": assistant})
messages.append({"role": "user", "content": user_message})
prompt = tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True,
)
return prompt
def chat_logic(message, history, system_prompt, context, max_tokens, temperature, top_p):
if not message: return "", history
if history is None: history = []
if not system_prompt: system_prompt = DEFAULT_SYSTEM_PROMPT
max_tokens = int(min(max_tokens, 1024))
prompt = build_prompt(system_prompt, context, history, message)
inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(DEVICE)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=max_tokens,
do_sample=True,
temperature=float(temperature),
top_p=float(top_p),
pad_token_id=tokenizer.eos_token_id,
)
decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
if decoded.startswith(prompt):
answer = decoded[len(prompt):].strip()
else:
answer = decoded.split("assistant\n")[-1].strip()
history.append((message, answer))
return "", history
# ==========================================
# 3. UI DESIGN (FINAL FIX)
# ==========================================
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;600;700&display=swap');
/* --- GLOBAL RESET --- */
:root, .dark, body, .gradio-container {
--body-text-color: #000000 !important;
--body-background-fill: transparent !important;
--background-fill-primary: transparent !important;
--background-fill-secondary: transparent !important;
--block-background-fill: transparent !important;
--block-border-color: transparent !important;
--input-background-fill: #ffffff !important;
/* COLORS */
--primary-500: #000000 !important;
--primary-600: #000000 !important;
--secondary-500: #000000 !important;
}
body {
font-family: 'Inter', sans-serif !important;
background: radial-gradient(circle at 10% 10%, #ffeef8 0%, #edf6ff 50%, #f7f0ff 90%) !important;
color: #000000 !important;
margin: 0; padding: 0;
min-height: 100vh;
}
/* HIDE GRADIO FOOTER (CSS is enough) */
footer {
display: none !important;
visibility: hidden !important;
}
h1 { font-weight: 700 !important; color: #000000 !important; letter-spacing: -0.02em; }
p, span, label { color: #374151 !important; }
/* --- LEFT PANEL --- */
.config-card {
background: rgba(255, 255, 255, 0.4) !important;
border: 1px solid rgba(255, 255, 255, 0.5) !important;
backdrop-filter: blur(10px);
border-radius: 24px;
padding: 24px !important;
}
.config-card div, .config-card .gr-form, .config-card .gr-box,
.config-card .gr-group, .config-card .gr-block, .config-card .gr-panel {
background: transparent !important;
border: none !important;
box-shadow: none !important;
}
.config-card label span,
.config-card .label-wrap span,
.config-card .block-title {
background-color: transparent !important;
background: transparent !important;
color: #000000 !important;
border: none !important;
font-weight: 700 !important;
font-size: 0.85rem !important;
margin-bottom: 4px !important;
padding: 0 !important;
}
.config-card textarea,
.config-card input[type="text"],
.config-card input[type="number"] {
background-color: #ffffff !important;
color: #000000 !important;
border: 1px solid #e5e7eb !important;
border-radius: 12px !important;
box-shadow: 0 1px 2px rgba(0,0,0,0.05) !important;
}
input[type=range] {
accent-color: #000000 !important;
filter: grayscale(100%) brightness(0) !important;
}
/* --- RIGHT PANEL: CHAT --- */
.chat-window {
background: rgba(255, 255, 255, 0.95) !important;
border: 1px solid rgba(255, 255, 255, 0.9);
backdrop-filter: blur(20px);
border-radius: 24px;
box-shadow: 0 20px 50px rgba(0,0,0,0.05);
overflow: hidden;
padding: 0 !important;
}
.window-header {
display: flex; gap: 8px; align-items: center;
padding: 18px 24px 10px 24px;
}
.dot { width: 12px; height: 12px; border-radius: 50%; }
.red { background-color: #ff5f56; }
.yellow { background-color: #ffbd2e; }
.green { background-color: #27c93f; }
#chatbot {
background: transparent !important;
border: none !important;
height: 550px !important;
}
/* --- CHAT BUBBLES --- */
.message.user {
background-color: #000000 !important;
color: #ffffff !important;
border-radius: 18px 18px 4px 18px !important;
font-weight: 500;
}
.message.user p, .message.user span, .message.user div {
color: #ffffff !important;
}
.message.bot {
background-color: #ffffff !important;
color: #000000 !important;
border: 1px solid #e5e7eb !important;
border-radius: 18px 18px 18px 4px !important;
box-shadow: 0 2px 5px rgba(0,0,0,0.02);
}
/* --- INPUT BOX --- */
.chat-input-area {
padding: 16px 24px 24px 24px;
display: flex; gap: 10px; align-items: flex-end;
}
#msg-box textarea {
background: #ffffff !important;
padding: 12px 18px !important;
border: 1px solid #e5e7eb !important;
color: #000000 !important;
max-height: 120px !important;
overflow-y: auto !important;
border-radius: 24px !important;
}
#send-btn {
background: transparent !important;
color: #000000 !important;
border: none !important;
box-shadow: none !important;
font-size: 1.6rem;
padding: 0 10px !important;
width: auto !important;
margin-bottom: 8px;
}
#send-btn:hover {
transform: translateX(3px);
color: #333;
}
"""
WINDOW_HEADER = """
<div class="window-header">
<div class="dot red"></div>
<div class="dot yellow"></div>
<div class="dot green"></div>
<span style="margin-left: 12px; font-size: 0.85rem; color: #6b7280; font-weight: 500;">Iris Assistant</span>
</div>
"""
with gr.Blocks(theme=gr.themes.Soft(), css=CSS, title="Custom AI") as demo:
with gr.Row(elem_id="main-container"):
# --- LEFT PANEL ---
with gr.Column(scale=4):
gr.Markdown(
"""
# MX Iris AI Chatbot
"""
)
with gr.Group(elem_classes="config-card"):
gr.Markdown("### ⚙️ Settings")
context_input = gr.Textbox(
label="Context (Data)",
placeholder="Paste context, notes, or data here...",
lines=5
)
system_prompt_input = gr.Textbox(
label="System Prompt",
value=DEFAULT_SYSTEM_PROMPT,
lines=2
)
with gr.Row():
temp_slider = gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature")
tokens_slider = gr.Slider(64, 1024, value=256, step=64, label="Max Tokens")
top_p_slider = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-P")
# --- RIGHT PANEL ---
with gr.Column(scale=6):
with gr.Column(elem_classes="chat-window"):
gr.HTML(WINDOW_HEADER)
chatbot = gr.Chatbot(
elem_id="chatbot",
show_label=False,
avatar_images=(None, "https://cdn-icons-png.flaticon.com/512/4712/4712027.png"),
height=550,
layout="bubble"
)
with gr.Row(elem_classes="chat-input-area"):
msg = gr.Textbox(
show_label=False,
placeholder="Ask a question...",
scale=9,
elem_id="msg-box",
container=False,
autofocus=True,
lines=1
)
send_btn = gr.Button("➝", scale=1, elem_id="send-btn")
inputs_list = [msg, chatbot, system_prompt_input, context_input, tokens_slider, temp_slider, top_p_slider]
msg.submit(chat_logic, inputs_list, [msg, chatbot])
send_btn.click(chat_logic, inputs_list, [msg, chatbot])
if __name__ == "__main__":
# Removed incorrect show_footer argument. CSS handles footer hiding.
demo.launch(show_api=False)