Update app.py

app.py

@@ -2,20 +2,12 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Load model and tokenizer
-model_name = "inceptionai/jais-13b"
+# Load JAIS-13B model and tokenizer
+model_name = "InceptionAI/jais-13b"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
 
-# Define chatbot function
-def chat_with_jais(prompt):
+def generate_response(prompt, max_length=512, temperature=0.7):
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-    outputs = model.generate(**inputs, max_length=512)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# Gradio Interface
-interface = gr.Interface(fn=chat_with_jais, inputs="text", outputs="text", title="JAIS-13B Chatbot")
-
-# Launch the app
-if __name__ == "__main__":
-    interface.launch()
+    with torch.no_grad():
+        output = model