code_compiler / translator.py
MLDeveloper's picture
Update translator.py
4dfb3f4 verified
raw
history blame
1.58 kB
import streamlit as st
import requests
import os
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}:\n\n{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("🔄 Code Translator (Python, Java, C++, C)")
st.write("Translate 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 here:", height=200)
if st.button("Translate"):
if code_input.strip():
with st.spinner("Translating..."):
translated_code = translate_code(code_input, source_lang, target_lang)
st.subheader("Translated Code:")
st.code(translated_code, language=target_lang.lower())
else:
st.warning("⚠️ Please enter some code before translating.")