sarvam_assistant.py

import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

st.set_page_config(page_title="Sarvam AI Assistant", layout="centered")

@st.cache_resource
def load_model():
    model_id = "sarvamai/sarvam-m"
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto"
    )
    return tokenizer, model

tokenizer, model = load_model()

st.title("🧠 Sarvam-M AI Assistant 🇮🇳")
prompt = st.text_area("💬 Ask your question (in English or Indian languages):", height=150)

if st.button("🧾 Get Answer"):
    if prompt.strip() == "":
        st.warning("Please enter a prompt.")
    else:
        with st.spinner("Thinking..."):
            messages = [{"role": "user", "content": prompt}]
            text = tokenizer.apply_chat_template(
                messages, tokenize=False, enable_thinking=True
            )
            inputs = tokenizer([text], return_tensors="pt").to(model.device)
            output_ids = model.generate(**inputs, max_new_tokens=512)
            output_text = tokenizer.decode(output_ids[0][len(inputs["input_ids"][0]):])
            content = output_text.split("</think>")[-1].strip() if "</think>" in output_text else output_text
            st.success("✍️ Answer:")
            st.write(content)