Update app.py

app.py

@@ -1,21 +1,24 @@
 import gradio as gr
 import torch
 from transformers import pipeline
-import requests
 import re
 import os
 from huggingface_hub import login
+from gradio_client import Client
 
 # Authenticate with Hugging Face
 if "HF_TOKEN" in os.environ:
     login(token=os.environ["HF_TOKEN"])
 
-# Global variable to store the Atlas-Chat model
+# Global variables
 atlas_pipe = None
+transliteration_client = None
 
-def load_atlas_model():
-    """Load only the Atlas-Chat model locally"""
-    global atlas_pipe
+def load_models():
+    """Load Atlas-Chat model and setup transliteration client"""
+    global atlas_pipe, transliteration_client
+    
+    # Load Atlas-Chat model locally
     if atlas_pipe is None:
         print("🏔️ Loading Atlas-Chat-2B model...")
         atlas_pipe = pipeline(
@@ -25,7 +28,19 @@ def load_atlas_model():
             device="cuda" if torch.cuda.is_available() else "cpu"
         )
         print("✅ Atlas-Chat model loaded!")
-    return atlas_pipe
+    
+    # Setup transliteration client
+    if transliteration_client is None:
+        try:
+            # REPLACE WITH YOUR ACTUAL HELPER SPACE URL
+            print("🔗 Connecting to transliteration service...")
+            transliteration_client = Client("YOUR-USERNAME/arabizi-transliteration-helper")
+            print("✅ Transliteration client connected!")
+        except Exception as e:
+            print(f"❌ Failed to connect to transliteration service: {e}")
+            transliteration_client = None
+    
+    return atlas_pipe, transliteration_client
 
 def detect_arabizi(text):
     """
@@ -84,65 +99,29 @@ def detect_arabizi(text):
     
     return False
 
-def arabizi_to_arabic_api(arabizi_text):
+def arabizi_to_arabic_client(arabizi_text):
     """
-    Convert Arabizi text to Arabic using Hugging Face Inference API
+    Convert Arabizi text to Arabic using the helper Space
     """
     try:
-        # Check if HF_TOKEN is available
-        if "HF_TOKEN" not in os.environ:
-            print("❌ HF_TOKEN not found, falling back to original text")
+        _, client = load_models()
+        
+        if client is None:
+            print("❌ Transliteration client not available, using fallback")
             return arabizi_text
         
-        API_URL = "https://api-inference.huggingface.co/models/atlasia/Transliteration-Moroccan-Darija"
-        headers = {"Authorization": f"Bearer {os.environ['HF_TOKEN']}"}
+        # Call the helper space
+        result = client.predict(arabizi_text, api_name="/predict")
         
-        # Prepare the payload
-        payload = {
-            "inputs": arabizi_text,
-            "parameters": {
-                "max_length": 512,
-                "num_beams": 4,
-                "early_stopping": True
-            }
-        }
+        # Check if result is an error
+        if isinstance(result, str) and result.startswith("Error:"):
+            print(f"❌ Transliteration service error: {result}")
+            return arabizi_text
         
-        # Make API request with timeout
-        response = requests.post(API_URL, headers=headers, json=payload, timeout=30)
+        return result.strip() if result else arabizi_text
         
-        # Check if request was successful
-        if response.status_code == 200:
-            result = response.json()
-            
-            # Handle different response formats
-            if isinstance(result, list) and len(result) > 0:
-                if "generated_text" in result[0]:
-                    return result[0]["generated_text"].strip()
-                elif isinstance(result[0], str):
-                    return result[0].strip()
-            elif isinstance(result, dict) and "generated_text" in result:
-                return result["generated_text"].strip()
-            elif isinstance(result, str):
-                return result.strip()
-            else:
-                print(f"❌ Unexpected API response format: {result}")
-                return arabizi_text
-                
-        elif response.status_code == 503:
-            print("⏳ Model is loading, falling back to original text")
-            return arabizi_text
-        else:
-            print(f"❌ API error {response.status_code}: {response.text}")
-            return arabizi_text
-            
-    except requests.exceptions.Timeout:
-        print("⏰ API timeout, falling back to original text")
-        return arabizi_text
-    except requests.exceptions.RequestException as e:
-        print(f"❌ API request failed: {e}")
-        return arabizi_text
     except Exception as e:
-        print(f"❌ Unexpected error in API conversion: {e}")
+        print(f"❌ Error calling transliteration service: {e}")
         return arabizi_text
 
 def arabic_to_arabizi(arabic_text):
@@ -215,13 +194,13 @@ def arabic_to_arabizi(arabic_text):
     return result.strip()
 
 def chat_with_atlas(message, history):
-    """Generate response from Atlas-Chat model with API-powered Arabizi conversion"""
+    """Generate response from Atlas-Chat model with Space-to-Space Arabizi conversion"""
     if not message.strip():
         return "ahlan wa sahlan! kifash n9der n3awnek? / مرحبا! كيفاش نقدر نعاونك؟"
     
     try:
-        # Load Atlas-Chat model
-        atlas_model = load_atlas_model()
+        # Load models
+        atlas_model, _ = load_models()
         
         # Detect if input is Arabizi
         is_arabizi_input = detect_arabizi(message)
@@ -234,8 +213,8 @@ def chat_with_atlas(message, history):
         
         # Prepare input for the model
         if is_arabizi_input:
-            print("🔄 Converting Arabizi→Arabic via API...")
-            arabic_input = arabizi_to_arabic_api(message)
+            print("🔄 Converting Arabizi→Arabic via Helper Space...")
+            arabic_input = arabizi_to_arabic_client(message)
             print(f"✅ ARABIC: '{arabic_input}'")
             model_input = arabic_input
         else: