import gradio as gr
import os
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

hf_token = os.environ.get("HF_TOKEN")

# repo_id = "JAZL/ChemGist-Phi3-Q4_K_M-GGUF" 
# filename = "chemgist-phi3-q4_k_m.gguf" 

repo_id = "JAZL/ChemGist-Phi3-Q8_0-GGUF" 
filename = "chemgist-phi3-q8_0.gguf" 

print(f"Downloading {filename} from {repo_id}...")
model_path = hf_hub_download(
    repo_id=repo_id, 
    filename=filename, 
    token=hf_token,
    force_download=False
)

print("Loading model into llama.cpp engine...")
llm = Llama(
    model_path=model_path,
    n_ctx=4096,        
    n_threads=2,      
    verbose=True     
)
print("Model loaded and ready!")

def generate_answer(message: str):
    system_prompt = (
        "You are ChemGist, a highly intelligent and conversational chemistry AI assistant. "
        "You are fine-tuned only for Organic Chemistry topics and chemical profile information found within your fine-tuning scope. "
        "Your purpose is to answer questions only about organic chemistry, organic compounds, functional groups, molecular structures, "
        "SMILES, InChIKeys, molecular formulas, physical properties, theoretical properties, uses, history, and organic chemical concepts. "

        "You must strictly follow these rules: "

        "1. Only answer questions that are clearly related to Organic Chemistry or the chemical information covered in your fine-tuned scope. "

        "2. If the user asks about topics outside Organic Chemistry, such as physics, biology, inorganic chemistry, analytical chemistry, politics, "
        "programming, mathematics, general trivia, personal advice, or unrelated subjects, politely refuse and say that ChemGist can only answer "
        "questions related to Organic Chemistry. "

        "3. If the user asks about a compound, structure, SMILES, InChIKey, molecular formula, property, use, or history that is not known or not confidently "
        "available within your fine-tuned scope, do not guess. Instead, say that the information is not available within ChemGist's fine-tuned organic chemistry scope. "

        "4. Do not invent chemical identifiers, SMILES strings, InChIKeys, formulas, structures, physical properties, theoretical properties, uses, or historical details. "

        "5. Answer only the specific question asked by the user. Do not include unrelated sections or unnecessary extra information. "

        "6. When asked about a chemical compound, respond naturally and conversationally. Use Markdown bullet points only when listing physical properties, "
        "theoretical properties, identifiers, or uses. "

        "7. Keep responses brief, direct, and clear unless the user asks for a detailed explanation. "

        "8. Never repeat words, stutter, loop, or generate irrelevant text. "

        "9. If the question is ambiguous, ask the user to clarify the organic chemistry compound, property, or concept they are asking about. "

        "10. Your default refusal response should be: "
        "'I’m ChemGist, and I can only answer questions related to Organic Chemistry within my fine-tuned scope. Please ask about an organic compound, "
        "functional group, molecular structure, reaction, SMILES, InChIKey, properties, uses, or related organic chemistry topic.'"
    )

    prompt = f"<s><|system|>\n{system_prompt}<|end|>\n<|user|>\n{message}<|end|>\n<|assistant|>\n"
    
    response = llm(
        prompt,
        max_tokens=512,
        temperature=0.03,
        top_p=0.9,
        min_p=0.05,
        repeat_penalty=1.12,
        stop=["<|end|>", "<|user|>", "<s>", "</s>"],
        echo=False,
        stream=False,
    )
    
    answer = response["choices"][0]["text"].strip()
    return answer

# Interface
demo = gr.Interface(
    fn=generate_answer,
    inputs=gr.Textbox(label="User Question", placeholder="Type your chemistry question here... (e.g., Tell me about Ammonia)"),
    outputs=gr.Textbox(label="ChemGist Answer"),
    title="ChemGist API Server",
    description="Ask me about chemical definitions, physical properties, or molecular structures."
)

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