sarvamai / app.py
sagar007's picture
Update app.py
e2d8960 verified
raw
history blame
1.3 kB
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()