Spaces:
Sleeping
Sleeping
File size: 3,059 Bytes
15c7196 5333011 15c7196 5333011 15c7196 5333011 15c7196 |
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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import os
import torch
import torch.nn as nn
import streamlit as st
from pydantic import BaseModel
from fastapi import FastAPI, Request
from transformers import AutoTokenizer, AutoModel
from peft import PeftModel
# Get the token from environment variable (optional)
hf_token = os.environ.get("HF_TOKEN")
# Define model IDs
adapter_model_id = "seniormgt/arabicmgt-test"
base_model_id = "Alibaba-NLP/gte-multilingual-base"
# Define your model
class GTEClassifier(nn.Module):
def __init__(self, model_name=base_model_id):
super(GTEClassifier, self).__init__()
self.base_model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
self.config = self.base_model.config
self.pooler = nn.Linear(self.config.hidden_size, self.config.hidden_size)
self.pooler_activation = nn.Tanh()
self.dropout = nn.Dropout(0.0)
self.classifier = nn.Linear(self.config.hidden_size, 1)
self.loss_fn = nn.BCEWithLogitsLoss()
def forward(self, input_ids=None, attention_mask=None, inputs_embeds=None, labels=None, **kwargs):
if inputs_embeds is not None:
outputs = self.base_model(inputs_embeds=inputs_embeds, attention_mask=attention_mask)
else:
outputs = self.base_model(input_ids=input_ids, attention_mask=attention_mask)
pooled_output = outputs.last_hidden_state[:, 0, :]
pooled_output = self.pooler(pooled_output)
pooled_output = self.pooler_activation(pooled_output)
logits = self.classifier(self.dropout(pooled_output)).squeeze(-1)
loss = None
if labels is not None:
labels = labels.float()
loss = self.loss_fn(logits, labels)
return {"loss": loss, "logits": logits}
# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(adapter_model_id, token=hf_token, trust_remote_code=True)
base_model = GTEClassifier()
peft_model = PeftModel.from_pretrained(base_model, adapter_model_id, token=hf_token)
# peft_model.eval()
# Define prediction
def classify_text(text):
inputs = tokenizer(text, max_length=512, padding=True, return_attention_mask=True, return_tensors="pt", truncation=True)
input_ids = inputs['input_ids']
attention_mask = inputs['attention_mask']
with torch.no_grad():
outputs = peft_model(input_ids=input_ids, attention_mask=attention_mask)
logits = outputs["logits"]
probs = torch.sigmoid(logits).cpu().numpy().squeeze()
pred_label = int(probs >= 0.5)
return {"label": str(pred_label), "confidence": float(probs)}
# 🔹 Streamlit UI
st.title("Text Classification (MGT Detection)")
text = st.text_area("Enter text", height=150)
if st.button("Classify") and text.strip():
result = classify_text(text)
st.json(result)
# 🔹 FastAPI endpoint
app = FastAPI()
class Input(BaseModel):
data: list
@app.post("/predict")
async def predict(request: Request):
payload = await request.json()
text = payload["data"][0]["text"]
result = classify_text(text)
return {"data": [result]}
|