Update app.py

app.py

@@ -7,7 +7,9 @@ model_name = "maulanayyy/codet5_code_translation"  # Ganti dengan nama model yan
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-print("Model and tokenizer loaded successfully.")
+# Check if GPU is available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)  # Pindahkan model ke GPU jika tersedia
 
 # Function to perform inference
 def translate_code(input_code):
@@ -15,15 +17,12 @@ def translate_code(input_code):
         # Prepare the input text
         input_text = f"translate Java to C#: {input_code}"
         
-        # Check if GPU is available
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        
         # Tokenize the input
         input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)  # Pastikan input_ids ada di GPU
         
         # Generate the output
         with torch.no_grad():
-            outputs = model.generate(input_ids, max_length=512)
+            outputs = model.generate(input_ids, max_length=256)  # Kurangi max_length jika perlu
         
         # Decode the output
         translated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)