Update app.py

app.py

@@ -15,11 +15,26 @@ import subprocess
 import webbrowser
 import urllib.parse
 import http.client
+import torch
+import deepspeed  # Import DeepSpeed for ZeroGPU
 
 # Suppress warnings
-warnings.filterwarnings('ignore', category=UserWarning)
+warnings.filterwarnings('ignore', category=User Warning)
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
 
+def initialize_zero_gpu():
+    """Initialize Hugging Face ZeroGPU"""
+    try:
+        # Check if CUDA is available
+        if torch.cuda.is_available():
+            # Initialize DeepSpeed with Zero Optimization
+            deepspeed.init_distributed()
+            print("ZeroGPU initialized successfully.")
+        else:
+            print("CUDA is not available. Using CPU instead.")
+    except Exception as e:
+        print(f"Error initializing Hugging Face ZeroGPU: {str(e)}")
+
 def initialize_environment():
     """Initialize application environment and configurations"""
     # Create necessary directories
@@ -116,7 +131,8 @@ class GitHubBot:
 
     def __init__(self):
         self.github_api = None
-    def initialize_api(self, token: str):
+
+    def initialize _api(self, token: str):
         """Initialize GitHub API with token"""
         self.github_api = GitHubAPI(token)
 
@@ -362,6 +378,9 @@ def signal_handler(signum, frame):
     sys.exit(0)
 
 if __name__ == "__main__":
+    # Initialize ZeroGPU
+    initialize_zero_gpu()
+
     # Register cleanup handlers
     atexit.register(cleanup)
     signal.signal(signal.SIGINT, signal_handler)
@@ -394,116 +413,4 @@ if __name__ == "__main__":
         logger.error(f"Error launching application: {str(e)}")
         raise
     finally:
-        logger.info("Application shutdown")
-
-from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
-
-class HuggingFaceModel:
-    """Class to handle Hugging Face model loading and predictions"""
-
-    def __init__(self, model_name: str):
-        self.model_name = model_name
-        self.model = None
-        self.tokenizer = None
-        self.load_model()
-
-    def load_model(self):
-        """Load the Hugging Face model and tokenizer"""
-        try:
-            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-            self.model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
-            logger.info(f"Model {self.model_name} loaded successfully.")
-        except Exception as e:
-            logger.error(f"Error loading model {self.model_name}: {str(e)}")
-            raise
-
-    def predict(self, text: str) -> Dict:
-        """Make a prediction using the loaded model"""
-        try:
-            inputs = self.tokenizer(text, return_tensors="pt")
-            outputs = self.model(**inputs)
-            predictions = outputs.logits.argmax(dim=-1).item()
-            logger.info(f"Prediction made for input: {text} with result: {predictions}")
-            return {"prediction": predictions}
-        except Exception as e:
-            logger.error(f"Error making prediction: {str(e)}")
-            return {"error": str(e)}
-
-# Update the Gradio interface to include model loading and prediction
-def create_gradio_interface():
-    """Create and configure Gradio interface with custom styling"""
-    bot = GitHubBot()
-    hf_model = None  # Initialize Hugging Face model variable
-
-    with gr.Blocks(css=custom_css, theme=gr.themes.Base()) as demo:
-        # Existing Gradio components...
-
-        model_name = gr.Textbox(
-            label="Hugging Face Model Name",
-            placeholder="Enter the model name (e.g., 'distilbert-base-uncased')",
-            elem_classes="input input-bordered input-primary"
-        )
-
-        load_model_button = gr.Button(
-            "Load Model",
-            elem_classes="button button-primary"
-        )
-
-        prediction_text = gr.Textbox(
-            label="Input Text for Prediction",
-            placeholder="Enter text to classify...",
-            elem_classes="textarea textarea-primary"
-        )
-
-        predict_button = gr.Button(
-            "Make Prediction",
-            elem_classes="button button-success"
-        )
-
-        output_prediction = gr.Textbox(
-            label="Prediction Output",
-            interactive=False,
-            elem_classes="output-area"
-        )
-
-        # Load model handler
-        def load_model_handler(model_name_input):
-            nonlocal hf_model
-            try:
-                hf_model = HuggingFaceModel(model_name_input)
-                return f"Model {model_name_input} loaded successfully."
-            except Exception as e:
-                return f"Error loading model: {str(e)}"
-
-        # Prediction handler
-        def predict_handler(input_text):
-            if hf_model is None:
-                return "Model not loaded. Please load a model first."
-            result = hf_model.predict(input_text)
-            return result
-
-        # Connect components
-        load_model_button.click(
-            load_model_handler,
-            inputs=[model_name],
-            outputs=[output_prediction]
-        )
-
-        predict_button.click(
-            predict_handler,
-            inputs=[prediction_text],
-            outputs=[output_prediction]
-        )
-
-    return demo
-
-def initialize_zero_gpu():
-    """Initialize Hugging Face ZeroGPU"""
-    try:
-        import torch
-        torch.cuda.is_available()
-        print("Hugging Face ZeroGPU is not available. Using GPU instead.")
-    except Exception as e:
-        print(f"Error initializing Hugging Face ZeroGPU: {str(e)}")
-
-initialize_zero_gpu()
+        logger.info("Application shutdown")