Update app.py

app.py

@@ -1,9 +1,11 @@
 import streamlit as st
 import requests
+import re
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
 import pandas as pd
 from datasets import Dataset
+from huggingface_hub import hf_api
 
 # Title and description
 st.title("OSINT Tool 🏢")
@@ -11,7 +13,7 @@ st.markdown("""
     This tool performs **Open Source Intelligence (OSINT)** analysis on GitHub repositories and fetches titles from URLs.
     It also allows uploading datasets (CSV format) for fine-tuning models like **DistilBERT**.
     """)
-    
+
 # Sidebar for navigation
 st.sidebar.title("Navigation")
 app_mode = st.sidebar.radio("Choose the mode", ["GitHub Repository Analysis", "URL Title Fetcher", "Dataset Upload & Fine-Tuning"])
@@ -98,18 +100,30 @@ elif app_mode == "Dataset Upload & Fine-Tuning":
                     
                     tokenized_datasets = dataset.map(preprocess_function, batched=True)
                     
-                    # Training loop (example)
-                    train_args = {
-                        "output_dir": "./results",
-                        "num_train_epochs": 3,
-                        "per_device_train_batch_size": 16,
-                        "logging_dir": "./logs",
-                    }
+                    # Fine-tuning setup (using Hugging Face Trainer for a complete setup)
+                    from transformers import Trainer, TrainingArguments
+
+                    training_args = TrainingArguments(
+                        output_dir="./results",
+                        evaluation_strategy="epoch",
+                        learning_rate=2e-5,
+                        per_device_train_batch_size=16,
+                        per_device_eval_batch_size=16,
+                        num_train_epochs=3,
+                        weight_decay=0.01,
+                    )
                     
-                    # Fine-tuning logic (for demonstration purposes, actual fine-tuning will need Hugging Face Trainer)
-                    # model.train()
+                    trainer = Trainer(
+                        model=model,
+                        args=training_args,
+                        train_dataset=tokenized_datasets,
+                        eval_dataset=tokenized_datasets,
+                    )
                     
-                    st.success("Fine-tuning started (demo)!")
+                    # Train the model
+                    trainer.train()
+
+                    st.success("Fine-tuning completed successfully!")
                 except Exception as e:
                     st.error(f"Error during fine-tuning: {e}")
             else: