Update app.py

app.py

@@ -1,60 +1,26 @@
 import streamlit as st
-import requests
-import os
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Ensure the Hugging Face API Token is available
-API_TOKEN = os.getenv("HF_API_TOKEN")
-if not API_TOKEN:
-    st.error("⚠️ API Token is missing! Please set HF_API_TOKEN as an environment variable.")
-    st.stop()
-
-# Define Hugging Face 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 between languages using Hugging Face API."""
-    prompt = f"""You are an expert AI in code translation. Convert the following {source_lang} code to {target_lang}:
-    
-    {source_lang} Code:
-    ```{source_lang.lower()}
-    {code_snippet}
-    ```
-
-    Now, provide the equivalent {target_lang} code:
     """
+    Translate code using CodeT5 model.
+    """
+    prompt = f"""Translate this {source_lang} code to {target_lang}:
 
-    try:
-        response = requests.post(API_URL, headers=HEADERS, json={"inputs": prompt})
-
-        # Handle API response
-        if response.status_code == 200:
-            result = response.json()
-            if isinstance(result, list) and result:
-                generated_text = result[0].get("generated_text", "")
-                translated_code = generated_text.split(f"```{target_lang.lower()}")[-1].strip("```").strip()
-                return translated_code
-            else:
-                return "⚠️ Error: Unexpected API response format."
-
-        elif response.status_code == 400:
-            return "⚠️ Error: Bad request. Please check your input."
-
-        elif response.status_code == 401:
-            return "⚠️ Error: Unauthorized. Check your API token."
-
-        elif response.status_code == 403:
-            return "⚠️ Error: Access Forbidden. You may need special model access."
-
-        elif response.status_code == 503:
-            return "⚠️ Error: Model is loading. Please wait a moment and try again."
-
-        else:
-            return f"⚠️ API Error {response.status_code}: {response.text}"
-
-    except requests.exceptions.RequestException as e:
-        return f"⚠️ Network Error: {str(e)}"
+{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")
@@ -74,4 +40,4 @@ if st.button("Translate"):
         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())
+            st.code(translated_code, language=target_lang.lower())