Update main.py

main.py

@@ -3,17 +3,16 @@ from dotenv import load_dotenv
 from transformers import TFBertForSequenceClassification, BertTokenizerFast
 import tensorflow as tf
 
-# Directly specify model and API key
-MODEL_NAME = "Erfan11/Neuracraft"
-API_KEY = "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"
+# Load environment variables from .env file
+load_dotenv()
 
-def load_model(Erfan11/Neuracraft):
+def load_model(model_name):
     # Load the TensorFlow model from Hugging Face Hub
-    model = TFBertForSequenceClassification.from_pretrained(model_name, use_auth_token=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)
+    model = TFBertForSequenceClassification.from_pretrained(model_name, use_auth_token=os.getenv('API_KEY'))
     return model
 
-def load_tokenizer(Erfan11/Neuracraft):
-    tokenizer = BertTokenizerFast.from_pretrained(model_name, use_auth_token=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)
+def load_tokenizer(model_name):
+    tokenizer = BertTokenizerFast.from_pretrained(model_name, use_auth_token=os.getenv('API_KEY'))
     return tokenizer
 
 def predict(text, model, tokenizer):
@@ -22,8 +21,8 @@ def predict(text, model, tokenizer):
     return outputs
 
 def main():
-    model_name = Erfan11/Neuracraft
-    model = load_model(Erfan11/Neuracraft)
+    model_name = os.getenv('MODEL_PATH')
+    model = load_model(model_name)
     tokenizer = load_tokenizer(model_name)
     # Example usage
     text = "Sample input text"