import gradio as gr
from transformers import pipeline

# Initialize the model
model_name = "sarvamai/sarvam-2b-v0.5"
pipe = pipeline("text-generation", model=model_name, device=-1)  # Set device to -1 for CPU

# Supported languages
LANGUAGES = ["English", "Bengali", "Gujarati", "Hindi", "Kannada", "Malayalam", "Marathi", "Oriya", "Punjabi", "Tamil", "Telugu"]

def chatbot(message, history, language):
    # Prepare the prompt
    prompt = f"Conversation in {language}:\n"
    for human, ai in history:
        prompt += f"Human: {human}\nAI: {ai}\n"
    prompt += f"Human: {message}\nAI:"

    # Generate response
    response = pipe(prompt, max_new_tokens=100, temperature=0.7, repetition_penalty=1.1)[0]['generated_text']
    
    # Extract only the AI's response
    ai_response = response.split("AI:")[-1].strip()
    
    return ai_response

# Create the Gradio interface
iface = gr.ChatInterface(
    chatbot,
    additional_inputs=[
        gr.Dropdown(choices=LANGUAGES, label="Select Language", value="English")
    ],
    title="Multilingual Indian Chatbot",
    description="Chat in multiple Indian languages using the sarvam-2b model.",
    examples=[
        ["नमस्ते, आप कैसे हैं?", "Hindi"]
    ],
    theme="soft"
)

# Launch the interface
iface.launch()