Textwizai / app.py
Erfan11's picture
Update app.py
da70a87 verified
raw
history blame
1.21 kB
import os
from transformers import TFBertForSequenceClassification, BertTokenizerFast
def load_model(model_name):
try:
# Load TensorFlow model first
model = TFBertForSequenceClassification.from_pretrained(model_name, use_auth_token="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")
except OSError:
# Fallback to PyTorch model if TensorFlow fails
model = TFBertForSequenceClassification.from_pretrained(model_name, use_auth_token="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", from_pt=True)
return model
def load_tokenizer(model_name):
tokenizer = BertTokenizerFast.from_pretrained(model_name, use_auth_token="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")
return tokenizer
def predict(text, model, tokenizer):
inputs = tokenizer(text, return_tensors="tf")
outputs = model(**inputs)
return outputs
def main():
# Replace 'Erfan11/Neuracraft' with the correct model path if necessary
model_name = "Erfan11/Neuracraft"
model = load_model(model_name)
tokenizer = load_tokenizer(model_name)
# Example prediction
text = "Sample input text"
result = predict(text, model, tokenizer)
print(result)
if __name__ == "__main__":
main()