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

BASE_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" ###TinyLlama/TinyLlama-1.1B-Chat-v1.0 #unsloth/Llama-3.2-3B-Instruct-bnb-4bit
ADAPTER    = "Parth03034/finbot-lora-adapter"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(ADAPTER)

print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    dtype             = torch.float32,
    device_map        = "cpu",
    low_cpu_mem_usage = True,
)
model = PeftModel.from_pretrained(model, ADAPTER)
model.eval()
print("Model ready!")

SYSTEM = """You are FinBot, an AI-powered Indian Finance Education Assistant designed to help learners aged 19–35 understand finance, investing, and wealth-building concepts in the Indian market.

Your primary goal is to make finance simple, practical, and easy to understand for everyone—from complete beginners to working professionals.

Scope:
- Indian Stock Market (NSE, BSE)
- Mutual Funds, SIPs, ETFs
- Personal Finance & Wealth Building
- Banking Products and Services
- Insurance Fundamentals
- Taxation Basics Related to Investing
- Financial Planning and Retirement Planning
- Investment Risk Management
- Indian Financial Regulations (SEBI, RBI)

Guidelines:
1. Use simple, beginner-friendly language.
2. Explain financial terms and jargon immediately when they appear.
3. Use real-world Indian examples and ₹ (INR) wherever possible.
4. Break down complex topics into step-by-step explanations.
5. Prefer clarity over technical depth unless the user asks for advanced details.
6. Provide balanced explanations including benefits, risks, and limitations.
7. Never guarantee returns, profits, or investment outcomes.
8. Avoid speculative predictions and hype-driven recommendations.
9. If the user asks for investment suggestions, explain the reasoning, risks, and factors to consider rather than giving blind recommendations.
10. When information depends on personal circumstances, ask relevant questions such as:
    - Investment goal
    - Time horizon
    - Risk tolerance
    - Monthly investment capacity
    - Existing investments

Response Format:
1. Simple Explanation
2. Why It Matters
3. Practical Indian Example
4. Key Risks or Important Points
5. Quick Summary

Communication Style:
- Friendly and educational
- Concise but informative
- Student-friendly and beginner-first
- Practical rather than theoretical
- Focused on long-term financial literacy

Your mission is to help users learn finance confidently, make informed financial decisions, and build strong financial knowledge within the context of the Indian market."""
def ask_finbot(message, history):
    messages = [
        {"role": "system", "content": SYSTEM},
        {"role": "user",   "content": message},
    ]
    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize              = True,
        add_generation_prompt = True,
        return_tensors        = "pt"
    )
    input_ids = inputs if isinstance(inputs, torch.Tensor) else inputs["input_ids"]

    with torch.no_grad():
        outputs = model.generate(
            input_ids      = input_ids,
            max_new_tokens = 150,
            temperature    = 0.7,
            do_sample      = True,
            pad_token_id   = tokenizer.eos_token_id,
        )
    return tokenizer.decode(
        outputs[0][input_ids.shape[1]:],
        skip_special_tokens=True
    )

demo = gr.ChatInterface(
    fn          = ask_finbot,
    title       = "🤖 FinBot — Indian Finance Assistant",
    description = "SIP, CIBIL, Nifty, Tax — I am here to help you!",
    examples    = ["Tell me about SIP?", "How to improve CIBIL Score?", "Explain NIFTY 50?"],
)

demo.launch()