Update app.py

app.py

@@ -1,11 +1,39 @@
-import os
 import gradio as gr
+from transformers import RobertaTokenizer, RobertaForSequenceClassification
+import torch
 
-# Load the Hugging Face access token from environment variables
-hf_token = os.getenv("HUGGING_FACE_TOKEN")
+# Load the model and tokenizer from the specified directory
+model_path = './finetuned_roberta'
+tokenizer = RobertaTokenizer.from_pretrained(model_path)
+model = RobertaForSequenceClassification.from_pretrained(model_path)
 
-if hf_token is None:
-    raise ValueError("Hugging Face access token not found in environment variables.")
+# Define the prediction function
+def classify_text(text):
+    # Tokenize the input text
+    inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=128)
+    
+    # Get the model's prediction
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    # Apply softmax to get probabilities
+    probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
+    
+    # Get the probability of the class '1'
+    prob_1 = probabilities[0][1].item()
+    
+    return {"Probability of being 1": prob_1}
 
-# Load the model using the access token
-gr.load("GoalZero/GoalZero-latest", hf_token=hf_token).launch()
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=classify_text,
+    inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
+    outputs="json",
+    title="Text Classification with RoBERTa",
+    description="Enter some text and get the probability of the text being classified as class 1.",
+    enable_queue=True,  # Ensure the API is enabled
+)
+
+# Launch the app
+if __name__ == "__main__":
+    iface.launch()