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

model_name = "Menlo/Lucy-128k"  # Replace with your actual model repo ID

print("Loading model and tokenizer (CPU-only)...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32)
model.to("cpu")
model.eval()
print("Model loaded.")

def generate_response(prompt, max_new_tokens=200):
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cpu")
    with torch.no_grad():
        output_ids = model.generate(
            input_ids,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            eos_token_id=tokenizer.eos_token_id,
        )
    generated_ids = output_ids[0][input_ids.shape[-1]:]
    response = tokenizer.decode(generated_ids, skip_special_tokens=True)
    return response.strip()

def chatbot_with_wiki(user_input, wiki_search_enabled):
    if wiki_search_enabled:
        try:
            summary = wikipedia.summary(user_input, sentences=2)
            combined_prompt = f"User query: {user_input}\nWikipedia summary: {summary}\nAnswer naturally:\n"
        except wikipedia.exceptions.DisambiguationError as e:
            return f"Your query is ambiguous. Possible options are: {e.options[:5]}"
        except wikipedia.exceptions.PageError:
            combined_prompt = f"User query: {user_input}\n"
        except Exception as e:
            return f"Wikipedia search error: {e}"
    else:
        combined_prompt = f"User query: {user_input}\n"

    response = generate_response(combined_prompt)
    return response

def respond(message, wiki_enabled, chat_history):
    if message is None or message.strip() == "":
        return chat_history or []
    chat_history = chat_history or []
    # Add user message
    chat_history.append({"role": "user", "content": message})
    # Generate bot response
    bot_response = chatbot_with_wiki(message, wiki_enabled)
    chat_history.append({"role": "assistant", "content": bot_response})
    return chat_history

with gr.Blocks() as demo:
    gr.Markdown("# AI Chatbot with Lucy-128k and Wikipedia Search (CPU mode)")

    chatbot = gr.Chatbot(type="messages", label="Chatbot")
    user_input = gr.Textbox(placeholder="Enter your message and press Enter", show_label=False)
    wiki_toggle = gr.Checkbox(label="Enable Wikipedia Search", value=True)

    user_input.submit(
        respond,
        inputs=[user_input, wiki_toggle, chatbot],
        outputs=chatbot,
    )

    clear_btn = gr.Button("Clear Chat")
    clear_btn.click(lambda: [], outputs=chatbot)

if __name__ == "__main__":
    demo.launch()