adjust to API funcs

app.py

@@ -1,37 +1,65 @@
 import gradio as gr
-from setfit import SetFitModel
+import requests
+import time
+import os
 
-# Load a pre-trained SetFit model (replace with your specific model)
-model = SetFitModel.from_pretrained("CundK/ePA-Classifier-Agreement")
+# Define the AgreementClassifier class
+class AgreementClassifier:
+    def __init__(self, api_token, api_url, backoff_factor=1):
+        self.api_token = api_token
+        self.api_url = api_url
+        self.headers = {"Authorization": f"Bearer {self.api_token}"}
+        self.backoff_factor = backoff_factor
 
-# Define the mapping of labels
-label_mapping = {0: "Ablehnung", 1: "Neutral", 2: "Befürwortung"}
+    def query(self, payload):
+        retries = 0
+        while True:
+            response = requests.post(self.api_url, headers=self.headers, json=payload)
+            if response.status_code == 503:
+                retries += 1
+                wait_time = self.backoff_factor * (2 ** (retries - 1))
+                print(f"503 Service Unavailable. Retrying in {wait_time} seconds...")
+                time.sleep(wait_time)
+            else:
+                response.raise_for_status()
+                return response.json()
+
+    def classify_text_topic(self, input_text):
+        result = self.query(
+            {
+                "inputs": input_text,
+                "parameters": {},
+            }
+        )
+        return result
+    
+# Initialize the classifier with API token and URL
+API_TOKEN = os.getenv("API_TOKEN")
+API_URL = os.getenv("API_URL")
+classifier = AgreementClassifier(API_TOKEN, API_URL)
 
 # Define the color mapping
 color_mapping = {
     "Ablehnung": "red",
     "Neutral": "yellow",
-    "Befürwortung": "green"
+    "Zustimmung": "green"
 }
 
-# Function to classify text and return the label with background color
+# Function to classify text using the API
 def classify_text(text):
-    # Get predictions from the model
-    predictions = model([text])
-    
-    # Convert the prediction tensor to an integer index
-    label_index = int(predictions[0].item())
+    # Get predictions from the classifier
+    predictions = classifier.classify_text_topic(text)
     
-    # Map the integer index to the corresponding label
-    label = label_mapping[label_index]
+    # Find the label with the highest score
+    predicted_label = max(predictions, key=lambda x: x['score'])['label']
     
     # Return the label with the appropriate background color
-    return f'<div style="background-color: {color_mapping[label]}; padding: 10px; border-radius: 5px;">{label}</div>'
+    return f'<div style="background-color: {color_mapping[predicted_label]}; padding: 10px; border-radius: 5px;">{predicted_label}</div>'
 
 # Create the Gradio interface using Blocks for custom layout
 with gr.Blocks() as interface:
-    gr.Markdown("# Text Classification with SetFit")
-    gr.Markdown("Enter a sentence and get it classified into 'Ablehnung', 'Neutral', or 'Befürwortung'.")
+    gr.Markdown("# Text Classification with API")
+    gr.Markdown("Enter a sentence and get it classified into 'Ablehnung', 'Neutral', or 'Zustimmung'.")
     
     # Input text box set for single-line input
     text_input = gr.Textbox(lines=1, placeholder="Enter your text here...")  # Single-line input