app.py

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

# Initialize model
model_path = hf_hub_download(
    repo_id="AstroMLab/AstroSage-8B-GGUF",
    filename="AstroSage-8B-Q8_0.gguf"
)

llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=4,
    chat_format="llama-3",
    seed=42,
    f16_kv=True,
    logits_all=False,
    use_mmap=True,
    use_gpu=True
)

# Placeholder responses for when context is empty
GREETING_MESSAGES = [
    "Greetings! I am AstroSage, your guide to the cosmos. What would you like to explore today?",
    "Welcome to our cosmic journey! I am AstroSage. How may I assist you in understanding the universe?",
    "AstroSage here. Ready to explore the mysteries of space and time. How may I be of assistance?",
    "The universe awaits! I'm AstroSage. What astronomical wonders shall we discuss?",
]

def get_random_greeting():
    return random.choice(GREETING_MESSAGES)

def respond_stream(message, history):
    if not message:  # Handle empty messages
        return

    system_message = "You are AstroSage, a highly knowledgeable AI assistant..." # ... (your system message)
    messages = [{"role": "system", "content": system_message}]

    # Format history correctly (especially important if you use clear)
    for user, assistant in history:
        messages.append({"role": "user", "content": user})
        if assistant:  # Check if assistant message exists
            messages.append({"role": "assistant", "content": assistant})

    messages.append({"role": "user", "content": message})

    try:
        response_content = ""
        for chunk in llm.create_chat_completion(
            messages=messages,
            max_tokens=512,
            temperature=0.7,
            top_p=0.9,
            stream=True
        ):
            delta = chunk["choices"][0]["delta"]
            if "content" in delta: # check if content exists in delta
                response_content += delta["content"]
                yield response_content # yield inside the loop for streaming
    except Exception as e:
        yield f"Error during generation: {e}"


# Display the welcome message as the first assistant message
initial_message = random.choice(GREETING_MESSAGES)
chatbot = gr.Chatbot(value=[[None, initial_message]])  # Set initial value here

with gr.Blocks() as demo:
    chatbot.render()
    clear = gr.Button("Clear")
    clear.click(lambda: None, None, chatbot, fn=lambda: []) 

    demo.queue().launch()