Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn.functional as F | |
| from torch_geometric.nn import GCNConv | |
| from sklearn.feature_extraction.text import TfidfVectorizer | |
| from sklearn.model_selection import train_test_split | |
| import json | |
| import pickle | |
| import os | |
| import re | |
| import string | |
| # Clean text function | |
| def clean_text(text): | |
| text = text.lower() | |
| text = re.sub(r'\[.*?\]', '', text) | |
| text = re.sub(r'https?://\S+|www\.\S+', '', text) | |
| text = re.sub(r'<.*?>+', '', text) | |
| text = re.sub(f'[{re.escape(string.punctuation)}]', '', text) | |
| text = re.sub(r'\n', ' ', text) | |
| text = re.sub(r'\w*\d\w*', '', text) | |
| text = re.sub(' +', ' ', text) | |
| return text.strip() | |
| # BiGNN model | |
| class BiGNN(torch.nn.Module): | |
| def __init__(self, input_dim, hidden_dim, output_dim): | |
| super(BiGNN, self).__init__() | |
| self.conv1 = GCNConv(input_dim, hidden_dim) | |
| self.conv2 = GCNConv(hidden_dim, output_dim) | |
| def forward(self, data): | |
| x, edge_index = data.x, data.edge_index | |
| x = F.relu(self.conv1(x, edge_index)) | |
| x = F.dropout(x, training=self.training) | |
| x = self.conv2(x, edge_index) | |
| return F.log_softmax(x, dim=1) | |
| # Load data | |
| with open("question_api/bgnn_training_dataset.json", "r", encoding="utf-8") as f: | |
| data = json.load(f) | |
| sentences = [clean_text(d["sentence"]) for d in data] | |
| labels = [d["label"] for d in data] | |
| # Vectorize | |
| vectorizer = TfidfVectorizer() | |
| X = vectorizer.fit_transform(sentences).toarray() | |
| y = torch.tensor(labels) | |
| # Build edge index (fake graph: each node is connected to itself) | |
| edge_index = torch.tensor([[0], [0]], dtype=torch.long) | |
| # Train/Test Split | |
| X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) | |
| # Convert to tensor | |
| X_train = torch.tensor(X_train, dtype=torch.float32) | |
| X_test = torch.tensor(X_test, dtype=torch.float32) | |
| # Create dummy PyG data objects | |
| train_data = type('Data', (object,), {})() | |
| train_data.x = X_train | |
| train_data.edge_index = edge_index | |
| test_data = type('Data', (object,), {})() | |
| test_data.x = X_test | |
| test_data.edge_index = edge_index | |
| # Train model | |
| model = BiGNN(input_dim=X.shape[1], hidden_dim=16, output_dim=3) | |
| optimizer = torch.optim.Adam(model.parameters(), lr=0.01) | |
| epochs = 50 | |
| for epoch in range(epochs): | |
| model.train() | |
| optimizer.zero_grad() | |
| out = model(train_data) | |
| loss = F.nll_loss(out, y_train) | |
| loss.backward() | |
| optimizer.step() | |
| print(f"Epoch {epoch+1}/{epochs}, Loss: {loss.item():.4f}") | |
| # Save model and vectorizer | |
| os.makedirs("question_api/artifacts", exist_ok=True) | |
| torch.save(model.state_dict(), "question_api/artifacts/biggn_model.pt") | |
| with open("question_api/artifacts/vectorizer.pkl", "wb") as f: | |
| pickle.dump(vectorizer, f) | |
| print("✅ Model and vectorizer saved.") | |