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

model_id = "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")

SYSTEM = "You are a friendly chatbot created by Anil Niraula. Your training data ended in late 2023. You do not have internet access. Be helpful and concise."

finance_keywords = ["stock", "invest", "portfolio", "allocation", "etf", "bond", "market", "dividend", "401k", "ira", "brokerage"]

@spaces.GPU
def generate(message, history):
    if message.strip().lower() in ["hi", "hello", "hey", "hi there", "hello there"]:
        return "Hi! I am a friendly chatbot created by Anil Niraula. I can assist with many subjects, but my training ended in late 2023 and I do not have access to the internet."

    if any(k in message.lower() for k in finance_keywords):
        system = SYSTEM + """
Focus on these facts:
- Asset allocation: mix of stocks/bonds/cash to balance risk and return.
- Taxable brokerage: capital gains and dividends taxed annually.
- Tax-advantaged (401k/IRA/Roth): taxes deferred or tax-free.
- S&P 500 long-term average: ~10% nominal / ~7% real annual return.
Avoid specific price predictions.
"""
    else:
        system = SYSTEM

    messages = [{"role": "system", "content": system}]
    for h in history:
        messages.append({"role": "user", "content": h[0]})
        messages.append({"role": "assistant", "content": h[1]})
    messages.append({"role": "user", "content": message})

    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=256)
    return tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)

gr.ChatInterface(generate, title="Chatbot by Anil Niraula").launch()