import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from peft import PeftModel

# 1) Base model & tokenizer
BASE_MODEL = "facebook/blenderbot-400M-distill"
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
base_model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL)

# 2) Attach your LoRA adapter
ADAPTER_REPO = "abinashnp/bayedger-chatbot"
model = PeftModel.from_pretrained(base_model, ADAPTER_REPO)

# 3) Build the text2text pipeline (no explicit device arg)
chatbot = pipeline(
    "text2text-generation",
    model=model,
    tokenizer=tokenizer,
    # device_map="auto"  # only if you use Accelerate; otherwise remove
)

# 4) System prompt (context) that always precedes user questions
SYSTEM_PROMPT = (
    "You are BayEdger’s AI assistant. You only answer FAQs about BayEdger’s "
    "services, pricing, and contact info. If you don’t know the answer, "
    "you must say exactly:\n"
    '"Sorry, I don’t have that info—please contact hello@bayedger.com."\n\n'
    "Here is what you should know about BayEdger:\n"
    "- AI‐powered websites and automation\n"
    "- Chatbots, email agents, process automation, analytics, content gen\n"
    "- Clear pricing tiers: Basic site ($400), Chatbot ($750+50/mo), Email ($1k+100/mo), etc.\n"
    "- Starter/Growth/Premium bundles\n"
    "- Contact: hello@bayedger.com, +1‐234‐559‐87994, 13 Madison St, NY\n\n"
)

def respond(query):
    # 5) Compose full prompt
    prompt = (
        SYSTEM_PROMPT
        f"question: {query}\n"
        "answer:"
    )

    # 6) Generate
    out = chatbot(
        prompt,
        max_new_tokens=128,
        do_sample=False,
        num_beams=2,
        early_stopping=True,
        pad_token_id=tokenizer.eos_token_id
    )[0]["generated_text"]

    # 7) Strip off everything up through our "answer:" token
    if "answer:" in out:
        reply = out.split("answer:", 1)[1].strip()
    else:
        reply = out.strip()

    # 8) Fallback: if the model didn’t produce anything substantial
    if len(reply) < 15 or "don't know" in reply.lower() or "sorry" in reply.lower():
        return "Sorry, I don’t have that info—please contact hello@bayedger.com."

    return reply

# 9) Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 BayEdger FAQ Chatbot")
    txt = gr.Textbox(placeholder="Ask me about BayEdger…", label="Your question")
    out = gr.Textbox(label="Answer")
    txt.submit(respond, txt, out)

demo.launch()