# app.py
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
import torch
import os

# Load FAISS vectorstore
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectorstore = FAISS.load_local("embeddings/am_index", embedding_model, allow_dangerous_deserialization=True)

# Load model
model_name = "HuggingFaceH4/zephyr-7b-beta"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# QA Function
def generate_answer(query, api_key):
    if api_key.strip() != "am123456":
        return "❌ Invalid API Key."

    docs = vectorstore.similarity_search(query, k=2)
    context = "\n".join([doc.page_content for doc in docs])

    prompt = f"""You are a domain expert in Additive Manufacturing.
Based on the context, answer the following question in a short and precise paragraph.

### Context:
{context}

### Question:
{query}

### Answer:
"""
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(device)
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=300, temperature=0.5, top_p=0.95, do_sample=False)

    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
    final_answer = decoded.split("### Answer:")[-1].strip()
    return final_answer

# Gradio UI
with gr.Blocks(css="""
    #main-title { width: 50%; }
    .image-class { float: right; width: 50px; margin-left: auto; }
""") as demo:
    with gr.Row():
        gr.Markdown("## 🧠 Additive Manufacturing LLM", elem_id="main-title")
        gr.Image("logo.png", elem_id="logo", show_label=False, show_download_button=False, height=50, container=False)

    gr.Markdown("Answer technical questions with a focused and clean response. No extra metadata.")

    query = gr.Textbox(label="Ask your Additive Manufacturing question")
    key = gr.Textbox(label="Enter API Key (e.g., am123456)")
    output = gr.Textbox(label="Answer")

    btn = gr.Button("Get Answer")
    btn.click(fn=generate_answer, inputs=[query, key], outputs=output)

    gr.Markdown("### 🎉 Hosted permanently on Hugging Face Spaces")

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