Update app.py

app.py

@@ -1,13 +1,16 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 
-# 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")
+# Load model and tokenizer
+model = AutoModelForCausalLM.from_pretrained("inceptionai/jais-13b", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("inceptionai/jais-13b", trust_remote_code=True)
 
-def generate_response(prompt, max_length=512, temperature=0.7):
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-    with torch.no_grad():
-        output = model
+def generate_text(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(**inputs)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# Create Gradio interface
+interface = gr.Interface(fn=generate_text, inputs="text", outputs="text")
+
+interface.launch()