Update app.py

app.py

@@ -1,87 +1,8 @@
-import os
-import logging
-from transformers import BertTokenizer, BertModel, BertPreTrainedModel, BertConfig
-import torch
-from torch import nn
 import gradio as gr
-import torch.nn.functional as F
 
-# Set environment variables
-os.environ['MPLCONFIGDIR'] = '/tmp/matplotlib-config'
-os.environ['XDG_CACHE_HOME'] = '/tmp/fontconfig-cache'
-os.environ['HF_HOME'] = '/tmp/hf_cache'
+def greet(name):
+    return f"Hello {name}!"
 
-# Ensure the cache directory exists
-if not os.path.exists('/tmp/hf_cache'):
-    os.makedirs('/tmp/hf_cache')
-
-logging.basicConfig(level=logging.INFO)
-
-class CustomBertForSequenceClassification(BertPreTrainedModel):
-    def __init__(self, config):
-        super().__init__(config)
-        self.bert = BertModel(config)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-        self.fc_positive = nn.Linear(config.hidden_size, 1)
-        self.fc_neutral = nn.Linear(config.hidden_size, 1)
-        self.fc_negative = nn.Linear(config.hidden_size, 1)
-        self.init_weights()
-
-    def forward(self, input_ids, attention_mask=None, token_type_ids=None):
-        outputs = self.bert(
-            input_ids,
-            attention_mask=attention_mask,
-            token_type_ids=token_type_ids,
-        )
-        pooled_output = outputs[1]
-        pooled_output = self.dropout(pooled_output)
-        logits_positive = self.fc_positive(pooled_output)
-        logits_neutral = self.fc_neutral(pooled_output)
-        logits_negative = self.fc_negative(pooled_output)
-        return logits_positive, logits_neutral, logits_negative
-
-# Load the tokenizer
-tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-
-# Load the model with the trained weights
-model = CustomBertForSequenceClassification(BertConfig.from_pretrained('bert-base-uncased'))
-model.load_state_dict(torch.load('bert_classifier.pth', map_location=torch.device('cpu')), strict=False)
-model.eval()
-
-# Define the class labels
-class_labels = ["Positive Surprise", "No Surprise", "Negative Surprise"]
-
-def classify_text(text):
-    logging.info("Classifying text...")
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
-    with torch.no_grad():
-        logits_positive, logits_neutral, logits_negative = model(**inputs)
-        prob_positive = torch.sigmoid(logits_positive).item()
-        prob_neutral = torch.sigmoid(logits_neutral).item()
-        prob_negative = torch.sigmoid(logits_negative).item()
-        prob_dict = {
-            "Positive Surprise": round(prob_positive, 4),
-            "No Surprise": round(prob_neutral, 4),
-            "Negative Surprise": round(prob_negative, 4)
-        }
-    logging.info(f"Probabilities: {prob_dict}")
-    return prob_dict
-
-# Health check function
-def health_check():
-    return "Healthy"
-
-# Set flagging_dir to a writable directory
-flagging_dir = "/tmp/flagged"
-os.makedirs(flagging_dir, exist_ok=True)
-
-iface = gr.Interface(
-    fn=classify_text,
-    inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
-    outputs=gr.Label(),
-    title="BERT Classifier for Surprises",
-    description="Enter text to get probabilities for positive surprise, negative surprise, or no surprise",
-    flagging_dir=flagging_dir
-)
+iface = gr.Interface(fn=greet, inputs="text", outputs="text")
 
 iface.launch(share=True)