import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load the model and tokenizer
model_name = "maulanayyy/codet5_code_translation-v3"  # Ganti dengan nama model yang benar
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# 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):
    try:
        # Prepare the input text
        input_text = f"translate Java to C#: {input_code}"
        
        # 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=256)  # Kurangi max_length jika perlu
        
        # Decode the output
        translated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return translated_code  # Kembalikan hasil akhir
        
    except Exception as e:
        print(f"Error during translation: {e}")
        return "An error occurred during translation."

# Create Gradio interface
demo = gr.Interface(fn=translate_code, inputs="text", outputs="text", title="Java to C# Code Translator", description="Enter Java code to translate it to C#.")

# Launch the interface
demo.launch(share=True)