updated with examples

app.py

@@ -4,44 +4,60 @@ 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"
+            formatted_text += "\n" + message["content"] + "\n"
         elif message["role"] == "user":
-            formatted_text += "<|user|>\n" + message["content"] + "\n"
+            formatted_text += "\n" + message["content"] + "\n"
         elif message["role"] == "assistant":
-            formatted_text += "<|assistant|>\n" + message["content"].strip() + eos + "\n"
+            formatted_text += "\n" + message["content"].strip() + eos + "\n"
         else:
             raise ValueError(
-                "Tulu chat template only supports 'system', 'user' and 'assistant' roles. Invalid role: {}.".format(
+                "Tulu chat template only supports 'system', 'user', and 'assistant' roles. Invalid role: {}.".format(
                     message["role"]
                 )
             )
-    formatted_text += "<|assistant|>\n"
+    formatted_text += "\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),
+def inference(input_prompts, model, tokenizer):
+    input_prompts = [
+        create_prompt_with_chat_format([{"role": "user", "content": input_prompt}], add_bos=False)
+        for input_prompt in input_prompts
+    ]
+
+    encodings = tokenizer(input_prompts, padding=True, return_tensors="pt")
+    encodings = encodings.to(device)
+
+    with torch.no_grad():
+        outputs = model.generate(encodings.input_ids, do_sample=False, max_length=250)
+
+    output_texts = tokenizer.batch_decode(outputs.detach(), skip_special_tokens=True)
+
+    input_prompts = [
+        tokenizer.decode(tokenizer.encode(input_prompt), skip_special_tokens=True) for input_prompt in input_prompts
+    ]
+    output_texts = [output_text[len(input_prompt) :] for input_prompt, output_text in zip(input_prompts, output_texts)]
+    return output_texts
+
+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)
+examples = [
+    ["मुझे अपने करियर के बारे में सुझाव दो", "मैं कैसे अध्ययन कर सकता हूँ?"],
+    ["कृपया मुझे एक कहानी सुनाएं", "ताजमहल के बारे में कुछ बताएं"],
+    ["मेरा नाम क्या है?", "आपका पसंदीदा फिल्म कौन सी है?"],
+]
+
+iface = gr.Chat(
+    model_fn=lambda input_prompts: inference(input_prompts, model, tokenizer),
+    inputs=["text"],
+    outputs="text",
+    examples=examples,
     title="Airavata Chatbot",
     theme="light",  # Optional: Set a light theme
 )