Spaces:
Sleeping
Sleeping
import streamlit as st | |
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
# Load CodeT5 model from Hugging Face | |
MODEL_NAME = "Salesforce/codet5-large" | |
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME) | |
def translate_code(code_snippet, source_lang, target_lang): | |
""" | |
Translate code using CodeT5 model. | |
""" | |
prompt = f"""Translate this {source_lang} code to {target_lang}: | |
{code_snippet}""" | |
# Tokenize and generate translation | |
inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512) | |
outputs = model.generate(**inputs, max_length=512) | |
# Decode the output | |
translated_code = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
return translated_code | |
# Streamlit UI | |
st.title("🔄 AI Code Translator") | |
st.write("Convert code between Python, Java, C++, and C.") | |
languages = ["Python", "Java", "C++", "C"] | |
source_lang = st.selectbox("Select Source Language", languages) | |
target_lang = st.selectbox("Select Target Language", languages) | |
code_input = st.text_area("Enter your code:", height=200) | |
if st.button("Translate"): | |
if source_lang == target_lang: | |
st.warning("⚠️ Source and target languages must be different!") | |
elif not code_input.strip(): | |
st.warning("⚠️ Please enter some code before translating.") | |
else: | |
with st.spinner("Translating... ⏳"): | |
translated_code = translate_code(code_input, source_lang, target_lang) | |
st.subheader(f"Translated {target_lang} Code:") | |
st.code(translated_code, language=target_lang.lower()) |