Create app.py

app.py

@@ -0,0 +1,77 @@
+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):
+        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 = self.loss_fn(logits, labels.float()) if labels is not None else None
+        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]}