Update translator.py

translator.py

@@ -1,42 +1,41 @@
-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 get_api_token():
+    """Fetch API token from environment variable."""
+    api_token = os.getenv("HF_API_TOKEN")
+    if not api_token:
+        raise ValueError("⚠️ API Token is missing! Please set HF_API_TOKEN as an environment variable.")
+    return api_token
 
-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.")
+API_TOKEN = get_api_token()
+MODEL_ID = "codellama/CodeLlama-7b-Instruct-hf"
+API_URL = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
+HEADERS = {"Authorization": f"Bearer {API_TOKEN}"}
 
-languages = ["Python", "Java", "C++", "C"]
+def translate_code(code_snippet, source_lang, target_lang):
+    """Translate code using the Hugging Face API."""
+    prompt = (
+        f"### Task: Convert {source_lang} code to {target_lang} code.\n\n"
+        f"### {source_lang} Code:\n\n{code_snippet}\n\n"
+        f"### {target_lang} Code:\n"
+    )
 
-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)
+    try:
+        response = requests.post(API_URL, headers=HEADERS, json={"inputs": prompt})
+        if response.status_code == 200:
+            result = response.json()
+            return result[0].get("generated_text", "⚠️ No translated code received.")
+        else:
+            return f"⚠️ API Error {response.status_code}: {response.text}"
+    except requests.exceptions.RequestException as e:
+        return f"⚠️ Network Error: {str(e)}"
 
-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.")
+# Example usage
+if __name__ == "__main__":
+    source_code = """
+def add(a, b):
+    return a + b
+"""
+    translated_code = translate_code(source_code, "Python", "Java")
+    print("Translated Java Code:\n", translated_code)