Update app.py

app.py

@@ -4,62 +4,58 @@ import os
 
 # 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()
+
 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"""You are a code translation AI. Your task is to translate {source_lang} code into {target_lang} accurately.
-    
-    Example Translation:
-    Python:
-    ```python
-    name = input("Enter your name: ")
-    print("Hello, " + name)
-    ```
-    
-    Java:
-    ```java
-    import java.util.Scanner;
-    public class Main {{
-        public static void main(String[] args) {{
-            Scanner scanner = new Scanner(System.in);
-            System.out.print("Enter your name: ");
-            String name = scanner.nextLine();
-            System.out.println("Hello, " + name);
-        }}
-    }}
-    ```
-
-    Now translate the following {source_lang} code to {target_lang}:
+    prompt = f"""
+    Translate the following {source_lang} code to {target_lang}:
     
     {code_snippet}
-
+    
     Translated {target_lang} Code:
     """
 
     response = requests.post(API_URL, headers=HEADERS, json={
-        "inputs": prompt,
-        "parameters": {
-            "max_new_tokens": 200,
-            "temperature": 0.2,
-            "top_k": 50,
-            "stop": ["\n\n", "#", "//", "'''"]
-        }
+        "inputs": prompt
     })
 
     if response.status_code == 200:
-        generated_text = response.json()[0]["generated_text"]
-
-        # Extract only the translated code
-        translated_code = generated_text.split(f"Translated {target_lang} Code:\n")[-1].strip()
-        
-        # Clean output by removing unnecessary text
-        translated_code = translated_code.replace("```", "").strip()
-        return translated_code
+        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."
+
     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")