Spaces:
Running
Running
File size: 2,297 Bytes
c3964c3 | 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 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | # models.py
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from PIL import Image
from torchvision import transforms
from torchvision.models import efficientnet_b0
# -------------------
# Текстовые модели
# -------------------
class TextModels:
def __init__(self):
# модель 1
self.tokenizer1 = AutoTokenizer.from_pretrained("VSAsteroid/ai-text-detector-hc3")
self.model1 = AutoModelForSequenceClassification.from_pretrained("VSAsteroid/ai-text-detector-hc3")
self.model1.eval()
# модель 2 (для ensemble)
self.tokenizer2 = AutoTokenizer.from_pretrained("silentone0725/text-detector-model-v2")
self.model2 = AutoModelForSequenceClassification.from_pretrained("silentone0725/text-detector-model-v2")
self.model2.eval()
def predict(self, text):
# модель 1
inputs1 = self.tokenizer1(text, return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
outputs1 = self.model1(**inputs1)
prob1 = torch.softmax(outputs1.logits, dim=1)[0][1].item() # AI probability
# модель 2
inputs2 = self.tokenizer2(text, return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
outputs2 = self.model2(**inputs2)
prob2 = torch.softmax(outputs2.logits, dim=1)[0][1].item() # AI probability
# ensemble
return (prob1 + prob2) / 2
# -------------------
# Изображения
# -------------------
class ImageModel:
def __init__(self):
self.model = efficientnet_b0(pretrained=True)
self.model.eval()
# трансформация для EfficientNet
self.transform = transforms.Compose([
transforms.Resize((224,224)),
transforms.ToTensor(),
transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
])
def predict(self, image: Image.Image):
img = self.transform(image).unsqueeze(0) # batch 1
with torch.no_grad():
outputs = self.model(img)
# для MVP используем sigmoid на один нейрон
prob = torch.sigmoid(outputs[:,0])[0].item()
return prob |