Update app.py

app.py

@@ -1,61 +1,33 @@
 import streamlit as st
 import requests
-import os
+import os  # Import os to access environment variables
 
 # Get API token from environment variable
-API_TOKEN = os.getenv("HF_API_TOKEN")  # Ensure you set this in your environment
-if not API_TOKEN:
-    st.error("⚠️ API Token is missing! Set HF_API_TOKEN in your environment.")
-    st.stop()
-
+API_TOKEN = os.getenv("HF_API_TOKEN")  # Fetch token securely
 MODEL_ID = "bigcode/starcoder"
 API_URL = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
 HEADERS = {"Authorization": f"Bearer {API_TOKEN}"}
 
 def translate_code(code_snippet, source_lang, target_lang):
-    """Translate code using Hugging Face API with improved prompt."""
-    prompt = f"""
-    Translate the following {source_lang} code to {target_lang}:
-    
-    {code_snippet}
-    
-    Translated {target_lang} Code:
-    """
+    """Translate code using Hugging Face API securely."""
+    prompt = f"Translate the following {source_lang} code to {target_lang}:\n\n{code_snippet}\n\nTranslated {target_lang} Code:\n"
 
     response = requests.post(API_URL, headers=HEADERS, json={
-        "inputs": prompt
+        "inputs": prompt,
+        "parameters": {
+            "max_new_tokens": 150,
+            "temperature": 0.2,
+            "top_k": 50,
+            "stop": ["\n\n", "#", "//", "'''"]
+        }
     })
 
     if response.status_code == 200:
-        response_json = response.json()
-
-        if isinstance(response_json, list) and len(response_json) > 0:
-            generated_text = response_json[0].get("generated_text", "").strip()
-            
-            if generated_text:
-                # Extract only the translated code
-                translated_code = generated_text.split(f"Translated {target_lang} Code:")[-1].strip()
-                return translated_code
-            else:
-                return "⚠️ Error: No translated text received."
-
-        else:
-            return "⚠️ Error: Unexpected response format."
-
-    elif response.status_code == 400:
-        return "⚠️ Error: Invalid request. Check input format."
-
-    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 access to this model."
-
-    elif response.status_code == 503:
-        return "⚠️ Error: Model is loading. Please wait and try again."
-
+        generated_text = response.json()[0]["generated_text"]
+        translated_code = generated_text.split(f"Translated {target_lang} Code:\n")[-1].strip()
+        return translated_code
     else:
-        return f"⚠️ Error {response.status_code}: {response.text}"
+        return f"Error: {response.status_code}, {response.text}"
 
 # Streamlit UI
 st.title("🔄 Code Translator using StarCoder")