Update app.py

app.py

@@ -1,49 +1,81 @@
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# Load model and tokenizer directly
+tokenizer = AutoTokenizer.from_pretrained("ai4bharat/Airavata")
+model = AutoModelForCausalLM.from_pretrained("ai4bharat/Airavata")
+
+def generate_response(prompt):
+    # Tokenize input prompt and generate response
+    inputs = tokenizer(prompt, return_tensors="pt", max_length=256, truncation=True)
+    outputs = model.generate(**inputs)
+    response = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+    
+    return response
+
+# Define Gradio interface
+iface = gr.Interface(
+    fn=generate_response,
+    inputs=gr.Textbox(),
+    outputs=gr.Textbox(),
+    live=True,
+    title="CAMAI",
+    description="Enter a prompt to generate text.",
+)
+
+# Launch Gradio interface
+iface.launch()
+
+
+
+
+
+# import torch
+# from transformers import AutoTokenizer, AutoModelForCausalLM
+# import gradio as gr
+
+# device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+# 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
+
+
+# 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.inference_mode():
+#         outputs = model.generate(encodings.input_ids, do_sample=False, max_new_tokens=250)
+
+#     output_texts = tokenizer.batch_decode(outputs.detach(), skip_special_tokens=True)
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-
-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
-
-
-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.inference_mode():
-        outputs = model.generate(encodings.input_ids, do_sample=False, max_new_tokens=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
+#     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"