File size: 1,079 Bytes
59dceb9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
# Define model paths based on uploaded repo files
MODEL_PATH = "trained_model2/distilroberta_model.pth"
TOKENIZER_DIR = "trained_model2/distilroberta_tokenizer"
# Load tokenizer
tokenizer_rl = AutoTokenizer.from_pretrained(TOKENIZER_DIR)
# Load model
model_rl = AutoModelForSequenceClassification.from_pretrained('distilroberta-base', num_labels=2)
model_rl.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu')))
model_rl.eval()
# RL model classification function
def classify_with_rl(text):
inputs = tokenizer_rl(text, return_tensors="pt", truncation=True, padding=True, max_length=128)
with torch.no_grad():
outputs = model_rl(**inputs)
probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
return {"spam_probability": max(0, min(1, float(probs[0][1])))}
# Create API
iface = gr.Interface(fn=classify_with_rl, inputs=gr.Textbox(), outputs="json")
# Launch API
if __name__ == "__main__":
iface.launch()
|