Update app.py

app.py

@@ -1,11 +1,8 @@
 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
+from src.github_analysis import analyze_github_repo
+from src.url_fetcher import fetch_url_title
+from src.fine_tune_helpers import fine_tune_model
 
 # Title and description
 st.title("OSINT Tool 🏢")
@@ -13,7 +10,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"])
@@ -26,22 +23,13 @@ if app_mode == "GitHub Repository Analysis":
     
     if st.button("Analyze Repository"):
         if repo_owner and repo_name:
-            try:
-                response = requests.get(f"https://api.github.com/repos/{repo_owner}/{repo_name}")
-                data = response.json()
-                
-                if response.status_code == 200:
-                    st.subheader("Repository Details")
-                    st.write(f"**Name**: {data['name']}")
-                    st.write(f"**Owner**: {data['owner']['login']}")
-                    st.write(f"**Stars**: {data['stargazers_count']}")
-                    st.write(f"**Forks**: {data['forks_count']}")
-                    st.write(f"**Language**: {data['language']}")
-                    st.write(f"**Description**: {data['description']}")
-                else:
-                    st.error(f"Error: {data.get('message', 'Something went wrong with the request')}")
-            except Exception as e:
-                st.error(f"Error occurred: {e}")
+            repo_data = analyze_github_repo(repo_owner, repo_name)
+            if repo_data:
+                st.subheader("Repository Details")
+                for key, value in repo_data.items():
+                    st.write(f"**{key}**: {value}")
+            else:
+                st.error("Failed to retrieve repository details.")
         else:
             st.warning("Please enter both repository owner and name.")
 
@@ -52,20 +40,11 @@ elif app_mode == "URL Title Fetcher":
     
     if st.button("Fetch Title"):
         if url:
-            try:
-                response = requests.get(url)
-                if response.status_code == 200:
-                    # Try to extract the title from the HTML
-                    match = re.search('<title>(.*?)</title>', response.text)
-                    if match:
-                        title = match.group(1)
-                        st.write(f"**Page Title**: {title}")
-                    else:
-                        st.warning("Title tag not found in the page")
-                else:
-                    st.error(f"Failed to retrieve the page. Status code: {response.status_code}")
-            except Exception as e:
-                st.error(f"Error occurred: {e}")
+            title = fetch_url_title(url)
+            if title:
+                st.write(f"**Page Title**: {title}")
+            else:
+                st.error("Failed to retrieve the page title.")
         else:
             st.warning("Please enter a valid URL.")
 
@@ -76,58 +55,4 @@ elif app_mode == "Dataset Upload & Fine-Tuning":
     uploaded_file = st.file_uploader("Upload a CSV file for fine-tuning", type="csv")
     
     if uploaded_file is not None:
-        # Load the CSV into a pandas DataFrame
-        df = pd.read_csv(uploaded_file)
-        
-        # Display dataset preview
-        st.subheader("Dataset Preview")
-        st.write(df.head())
-
-        # Convert CSV to Hugging Face dataset format
-        dataset = Dataset.from_pandas(df)
-        
-        model_name = st.selectbox("Select model for fine-tuning", ["distilbert-base-uncased"])
-        
-        if st.button("Fine-tune Model"):
-            if model_name:
-                try:
-                    model = AutoModelForSequenceClassification.from_pretrained(model_name)
-                    tokenizer = AutoTokenizer.from_pretrained(model_name)
-                    
-                    # Prepare the dataset
-                    def preprocess_function(examples):
-                        return tokenizer(examples['text'], truncation=True, padding=True)
-                    
-                    tokenized_datasets = dataset.map(preprocess_function, batched=True)
-                    
-                    # 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,
-                    )
-                    
-                    trainer = Trainer(
-                        model=model,
-                        args=training_args,
-                        train_dataset=tokenized_datasets,
-                        eval_dataset=tokenized_datasets,
-                    )
-                    
-                    # Train the model
-                    trainer.train()
-
-                    st.success("Fine-tuning completed successfully!")
-                except Exception as e:
-                    st.error(f"Error during fine-tuning: {e}")
-            else:
-                st.warning("Please select a model for fine-tuning.")
-    
-    else:
-        st.warning("Please upload a dataset.")
+        fine_tune_model(uploaded_file)