import streamlit as st
from transformers import pipeline

st.title("Kalki-2.5 AI Interface")

# 1. Load the model
@st.cache_resource
def load_model():
    return pipeline("text-generation", model="upmarking/kalki-2.5")

generator = load_model()

# 2. Define your custom System Prompt
system_prompt = "You are Kalki 2.5, an AI Model Developed by Upmakring Solutions in India."

# 3. UI for Input
user_input = st.text_area("Enter your prompt:")

if st.button("Generate"):
    if user_input:
        # Create the message list with your custom system prompt
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_input},
        ]
        
        with st.spinner("Generating..."):
            # The pipeline uses the model's chat template
            result = generator(messages, max_new_tokens=150)
            
            # Extract the assistant's response
            output_text = result[0]['generated_text'][-1]['content']
            
            st.subheader("Output:")
            st.write(output_text)
    else:
        st.warning("Please enter some text.")