added accelerate support

app.py

@@ -8,11 +8,8 @@ model_id = "BSC-LT/salamandraTA-2b"
 
 # Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id)
-
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
 # Move model to GPU if available
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
 languages = [ "Spanish", "Catalan", "English", "French", "German", "Italian", "Portuguese", "Euskera", "Galician",
              "Bulgarian", "Czech", "Lithuanian", "Croatian", "Dutch", "Romanian", "Danish", "Greek", "Finnish",
              "Hungarian", "Slovak", "Slovenian", "Estonian", "Polish", "Latvian", "Swedish", "Maltese",
@@ -21,14 +18,14 @@ languages = [ "Spanish", "Catalan", "English", "French", "German", "Italian", "P
 example_sentence = ["Ahir se'n va anar, va agafar les seves coses i es va posar a navegar."]
 
 
-@spaces.GPU
+@spaces.GPU(duration=120)
 def translate(input_text, source, target):
     sentences = input_text.split('\n')
     generated_text = []
     for sentence in sentences:
       prompt = f'[{source}] {sentence} \n[{target}]'
 
-      input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(device)
+      input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(model.device)
       output_ids = model.generate(input_ids, max_length=500, num_beams=5)
       input_length = input_ids.shape[1]
 

requirements.txt

@@ -2,4 +2,5 @@ torch
 transformers==4.46.2
 gradio==5.5.0
 protobuf==5.28.3
-sentencepiece==0.2.0
+sentencepiece==0.2.0
+accelerate==1.0.1