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

MODEL_ID = "arinbalyan/smollm2-finance"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    device_map="auto",
)

def generate(prompt, history):
    if not prompt.strip():
        return "", history
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=200,
            temperature=0.8,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )
    response = tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    history.append({"role": "user", "content": prompt})
    history.append({"role": "assistant", "content": response.strip()})
    return "", history

theme = (
    gr.themes.Soft(primary_hue="amber", neutral_hue="slate")
    .set(
        button_primary_background_fill_hover="#d97706",
        button_primary_text_color="white",
        input_background_fill="#f8fafc",
        input_border_color="#e2e8f0",
        input_border_color_focus="#d97706",
    )
)

css = """.gradio-container { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important; }
.main-header { text-align: center; padding: 2rem 1rem; background: linear-gradient(135deg, #d97706 0%, #d97706dd 100%); color: white; border-radius: 12px; margin-bottom: 1rem; }
.main-header h1 { font-size: 2rem; font-weight: 700; margin-bottom: 0.25rem; }
.main-header p { font-size: 1rem; opacity: 0.9; }
.footer { text-align: center; padding: 1rem; color: #64748b; font-size: 0.875rem; }"""

with gr.Blocks(title="Finance Advisor") as demo:
    with gr.Row():
        with gr.Column():
            gr.HTML(f"""<div class="main-header"><h1>💰 Finance Advisor</h1><p>Generate with AI</p></div>""")
    chatbot = gr.Chatbot(height=400)
    msg = gr.Textbox(label="Prompt", placeholder='e.g., What is compound interest?', lines=2)
    with gr.Row():
        submit = gr.Button("Generate", variant="primary", size="lg", scale=1)
        clear = gr.Button("Clear", size="lg", scale=0)
    gr.Examples(examples=["What is compound interest?", "How should I start investing in my 20s?", "Explain the difference between stocks and bonds", "What is an emergency fund?", "How do capital gains taxes work?"], inputs=msg, label="Try these")
    gr.HTML(f"""<div class="footer">SmolLM2-1.7B on finance-alpaca</div>""")
    submit.click(generate, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch(theme=theme, css=css)