Update main.py

main.py

@@ -1,29 +1,30 @@
 import os
-import tensorflow as tf
 from dotenv import load_dotenv
-from transformers import BertTokenizerFast
+from transformers import TFBertForSequenceClassification, BertTokenizerFast
+import tensorflow as tf
 
-# Load environment variables
-load_dotenv()
+# Directly specify model and API key
+MODEL_NAME = "Erfan11/Neuracraft"
+API_KEY = "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"
 
-def load_model(model_path):
-    # Load the TensorFlow model using from_tf=True
-    model = tf.keras.models.load_model(model_path)
+def load_model(model_name):
+    # Load the TensorFlow model from Hugging Face Hub
+    model = TFBertForSequenceClassification.from_pretrained(model_name, use_auth_token=API_KEY)
     return model
 
-def load_tokenizer(model_path):
-    tokenizer = BertTokenizerFast.from_pretrained(model_path)
+def load_tokenizer(model_name):
+    tokenizer = BertTokenizerFast.from_pretrained(model_name, use_auth_token=API_KEY)
     return tokenizer
 
 def predict(text, model, tokenizer):
     inputs = tokenizer(text, return_tensors="tf")
-    outputs = model(inputs)
+    outputs = model(**inputs)
     return outputs
 
 def main():
-    model_path = os.getenv('Erfan11/Neuracraft')
-    model = load_model(model_path)
-    tokenizer = load_tokenizer(model_path)
+    model_name = MODEL_NAME
+    model = load_model(model_name)
+    tokenizer = load_tokenizer(model_name)
     # Example usage
     text = "Sample input text"
     result = predict(text, model, tokenizer)