# app.py
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# -------------------------
# Model Setup
# -------------------------
model_name = "ibm-granite/granite-3.3-2b-instruct"  # Smaller variant for Spaces

print("Loading tokenizer and model...")
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Use device_map="auto" and 8-bit loading to reduce VRAM usage
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16,
    load_in_8bit=True,  # If bitsandbytes installed
)

print(f"Model loaded on device: {model.device}")

# -------------------------
# Chat function
# -------------------------
def chat_with_granite(prompt):
    try:
        # Tokenize input and move tensors to the model device
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        
        # Generate response
        outputs = model.generate(
            **inputs,
            max_new_tokens=150,
            do_sample=True,
            temperature=0.7,
            top_p=0.9
        )
        
        # Decode response
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response

    except Exception as e:
        # Return error instead of crashing
        return f"Error generating response: {e}"

# -------------------------
# Gradio Interface
# -------------------------
iface = gr.Interface(
    fn=chat_with_granite,
    inputs=gr.Textbox(lines=3, placeholder="Ask me something about finance..."),
    outputs=gr.Textbox(),
    title="Finance Chatbot",
    description="Ask your finance questions and Granite will answer!",
    allow_flagging="never"
)

# Launch interface
iface.launch()