Spaces:
Paused
Paused
File size: 2,004 Bytes
d9aa2f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load model and tokenizer
model_name = "ai4bharat/Airavata"
tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(device)
# Function for generating responses
def inference(message):
prompt = create_prompt_with_chat_format([{"role": "user", "content": message}], add_bos=False)
encoding = tokenizer(prompt, return_tensors="pt").to(device)
with torch.inference_mode():
output = model.generate(encoding.input_ids, do_sample=False, max_new_tokens=250)
return tokenizer.decode(output[0], skip_special_tokens=True)[len(message) :]
def create_prompt_with_chat_format(messages, bos="<s>", eos="</s>", add_bos=True):
formatted_text = ""
for message in messages:
if message["role"] == "system":
formatted_text += "<|system|>\n" + message["content"] + "\n"
elif message["role"] == "user":
formatted_text += "<|user|>\n" + message["content"] + "\n"
elif message["role"] == "assistant":
formatted_text += "<|assistant|>\n" + message["content"].strip() + eos + "\n"
else:
raise ValueError(
"Tulu chat template only supports 'system', 'user' and 'assistant' roles. Invalid role: {}.".format(
message["role"]
)
)
formatted_text += "<|assistant|>\n"
formatted_text = bos + formatted_text if add_bos else formatted_text
return formatted_text
# Create Gradio chat interface
iface = gr.ChatInterface(
fn=inference,
inputs=[gr.Textbox(lines=3, label="Ask me anything")],
outputs=gr.Textbox(label="Response", live=True),
title="Airavata Chatbot",
theme="light", # Optional: Set a light theme
)
iface.launch()
|