Update app.py

app.py

@@ -8,24 +8,21 @@ description = """ try this space to build downstream applications with [CohereFo
 
 checkpoint = "CohereForAI/aya-101"
 tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
-
-if device == "cuda":
-    model = model.half()
+model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint, torch_dtype=torch.float16, low_cpu_mem_usage=True, device_map="auto")
 
 @spaces.GPU
-def translate(text):
+def aya(text):
     """
     Translates the input text to English using the Aya model.
     Assumes the model can automatically detect the input language.
     """
+    model.to(device)
+
     inputs = tokenizer.encode(text, return_tensors="pt").to(device)
     
     outputs = model.generate(inputs, max_new_tokens=128)
     
+    
     translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return translation
 
@@ -33,12 +30,13 @@ def main():
     with gr.Blocks() as demo:
         gr.Markdown(title)
         gr.Markdown(description)
+        output_text = gr.Textbox(label="🌐Aya", interactive=False)
         with gr.Row():
-            input_text = gr.Textbox(label="Input Text")
-            output_text = gr.Textbox(label="🌐Aya", interactive=False)
-        input_text.change(fn=translate, inputs=input_text, outputs=output_text)
+            input_text = gr.Textbox(label="🗣️Input Text")
+            submit_button = gr.Button("Translate")  # Add a button to submit the input
+        submit_button.click(fn=aya, inputs=input_text, outputs=output_text)
     
     demo.launch()
 
 if __name__ == "__main__":
-    main()
+    main()