Update translator.py

translator.py

@@ -1,33 +1,42 @@
+import streamlit as st
 import requests
 import os
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Your Hugging Face API token (Replace 'your_token_here' with your actual token)
-
-API_TOKEN = os.getenv("HF_API_TOKEN") 
-
-# Define model and API endpoint
-MODEL_ID = "bigcode/starcoder"
-API_URL = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
-HEADERS = {"Authorization": f"Bearer {API_TOKEN}"}
+# 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 Hugging Face API (No local download needed).
+    Translate code using CodeT5 model.
     """
-    prompt = f"Translate the following {source_lang} code to {target_lang}:\n\n{code_snippet}\n\nTranslated {target_lang} Code:"
-
-    response = requests.post(API_URL, headers=HEADERS, json={"inputs": prompt})
-
-    if response.status_code == 200:
-        return response.json()[0]["generated_text"]
+    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:
-        return f"Error: {response.status_code}, {response.text}"
-
-# Example usage
-source_code = """
-def add(a, b):
-    return a + b
-"""
-translated_code = translate_code(source_code, "Python", "Java")
-print("Translated Java Code:\n", translated_code)
-
+        st.warning("⚠️ Please enter some code before translating.")