Update app.py

app.py

@@ -1,50 +1,18 @@
 import warnings
 import gradio as gr
-from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM
-import torch
+from proxy_model import RemoteModelProxy
 
 # Suppress the FutureWarning
 warnings.filterwarnings("ignore", category=FutureWarning, module="torch")
 
-# Load the model and tokenizer
-model_name = "deepseek-ai/DeepSeek-V3"
-revision = "4c1f24cc10a2a1894304c7ab52edd9710c047571"
-
-print(f"Loading tokenizer from {model_name}...")
-tokenizer = AutoTokenizer.from_pretrained(model_name, revision=revision, trust_remote_code=True)
-
-print(f"Loading configuration from {model_name}...")
-config = AutoConfig.from_pretrained(model_name, revision=revision, trust_remote_code=True)
-
-# Remove quantization configuration if it exists
-if hasattr(config, 'quantization_config'):
-    del config.quantization_config
-
-print(f"Loading model from {model_name}...")
-model = AutoModelForCausalLM.from_pretrained(model_name, config=config, revision=revision, trust_remote_code=True)
-
-# Check if the model loaded successfully
-if model is None:
-    print("Failed to load model. Exiting...")
-    exit(1)
-else:
-    print("Model loaded successfully.")
+# Load the model via the proxy
+model_proxy = RemoteModelProxy("deepseek-ai/DeepSeek-V3")
 
 # Define the text classification function
 def classify_text(text):
     try:
-        # Tokenize the input text
-        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-        # Pass the inputs to the model
-        logits = model(**inputs)
-        # Get the probabilities
-        probabilities = torch.softmax(logits, dim=-1).tolist()[0]
-        # Get the predicted class
-        predicted_class = torch.argmax(logits, dim=-1).item()
-        return {
-            "Predicted Class": predicted_class,
-            "Probabilities": probabilities
-        }
+        result = model_proxy.classify_text(text)
+        return result
     except Exception as e:
         print(f"Error during text classification: {e}")
         return {