Update app.py

app.py

@@ -2,17 +2,29 @@ import os
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
-import json
 
 os.environ["TRANSFORMERS_CACHE"] = "/tmp"
 
+# Suppress NVML warning
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+torch.cuda.init()
+
 MODEL_NAME = "s-nlp/roberta-base-formality-ranker"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
 
+# Move model to GPU if available
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = model.to(device)
 
+def calculate_formality_percentages(score):
+    # Convert score to grayscale percentage (0-100)
+    grayscale = int(score * 100)
+    # Use grayscale to determine formal/informal percentages
+    formal_percent = grayscale
+    informal_percent = 100 - grayscale
+    return formal_percent, informal_percent
+
 def predict_formality(text):
     # Tokenize input
     encoding = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
@@ -23,9 +35,8 @@ def predict_formality(text):
         logits = model(**encoding).logits
     score = logits.softmax(dim=1)[:, 1].item()
 
-    # Calculate percentages
-    formal_percent = round(score * 100)
-    informal_percent = 100 - formal_percent
+    # Calculate percentages using grayscale
+    formal_percent, informal_percent = calculate_formality_percentages(score)
     
     # Create response in the new format
     response = {
@@ -37,6 +48,7 @@ def predict_formality(text):
     
     return response
 
+# Create Gradio interface
 demo = gr.Interface(
     fn=predict_formality,
     inputs=gr.Textbox(label="Enter your text", lines=3),
@@ -52,4 +64,9 @@ demo = gr.Interface(
 
 # Launch the app
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True, 
+        enable_queue=True 
+    )