import gradio as gr
import os
from transformers import pipeline

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

pipe = pipeline(
    "text-generation",
    model="CoRover/BharatGPT-3B-Indic",
    device_map="auto",
    token=token
)

def chat_function(message, history):
    messages = [
        {"role": "system", "content": "You are a helpful assistant who responds in Hindi."},
        {"role": "user", "content": message}
    ]
    
    # Generate the response
    outputs = pipe(messages, max_new_tokens=256)
    
    # Depending on the pipeline version, it returns a list of dictionaries.
    # We extract the last message from the assistant.
    response = outputs[0]["generated_text"][-1]["content"]
    return response

demo = gr.ChatInterface(
    fn=chat_function, 
    title="BharatGPT Interactive Assistant",
    description="A live demo of CoRover's 3B Indic model. Ask me a question!"
)

demo.launch()