Spaces:
Running
Running
File size: 4,424 Bytes
24ba913 151aa67 a84c0b0 151aa67 a84c0b0 151aa67 35b0708 151aa67 |
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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
import streamlit as st
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
from typing import List, Dict
import os
# Initialize the Hugging Face pipeline (ensure to use a valid model)
model_name = "your_huggingface_model_name" # Ensure to use a valid model
tokenizer = AutoTokenizer.from_pretrained(model_name)
try:
model = AutoModelForCausalLM.from_pretrained(model_name)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
except Exception as e:
st.error(f"Error initializing the model '{model_name}': {e}")
# Function to generate OSINT results
def generate_osint_results(prompt: str, history: List[Dict[str, str]]) -> List[str]:
"""
Simulates generating OSINT-based results from the user's input.
Args:
prompt (str): The user's input to the simulator.
history (List[Dict]): The user's message history with timestamps.
Returns:
List[str]: A list of OSINT responses from the AI.
"""
# Validate inputs
if not prompt.strip():
return ["Error: Prompt cannot be empty."]
if not isinstance(history, list) or not all(isinstance(h, dict) for h in history):
return ["Error: History must be a list of dictionaries."]
# Prepare messages for the AI
messages = [{"role": "system", "content": f"Responding to OSINT prompt: {prompt}"}]
for val in history:
if "user" in val:
messages.append({"role": "user", "content": val["user"]})
if "assistant" in val:
messages.append({"role": "assistant", "content": val["assistant"]})
# Append the current user prompt
messages.append({"role": "user", "content": prompt})
# Generate a response using the Hugging Face model
try:
response = generator(messages[-1]["content"], max_length=100, num_return_sequences=1)
return [response[0]["generated_text"]]
except Exception as e:
return [f"Error generating response: {e}"]
# Function for fine-tuning the model with the uploaded dataset
def fine_tune_model(dataset: str) -> str:
"""
Fine-tunes the model using the uploaded dataset.
Args:
dataset (str): The path to the dataset for fine-tuning.
Returns:
str: A message indicating whether fine-tuning was successful or failed.
"""
try:
# Process the dataset (dummy processing for illustration)
with open(dataset, "r") as file:
data = file.readlines()
# Simulate fine-tuning with the provided dataset
# Here, you would use the data to fine-tune the model
# For this example, we're not actually fine-tuning the model.
model.save_pretrained("./fine_tuned_model")
return "Model fine-tuned successfully!"
except Exception as e:
return f"Error fine-tuning the model: {e}"
# Streamlit app interface
st.title("OSINT Tool")
st.write("This tool generates OSINT-based results and allows you to fine-tune the model with custom datasets.")
# User input for prompt and message history
prompt = st.text_area("Enter your OSINT prompt here...", placeholder="Type your prompt here...")
history = []
# Display message history
if "history" not in st.session_state:
st.session_state.history = []
# Display past conversation
st.write("### Message History:")
for msg in st.session_state.history:
st.write(f"**User**: {msg['user']}")
st.write(f"**Assistant**: {msg['assistant']}")
# Fine-tuning functionality
dataset_file = st.file_uploader("Upload a dataset for fine-tuning", type=["txt"])
if dataset_file is not None:
# Save the uploaded file
dataset_path = os.path.join("uploads", dataset_file.name)
with open(dataset_path, "wb") as f:
f.write(dataset_file.read())
# Fine-tune the model
fine_tuning_status = fine_tune_model(dataset_path)
st.success(fine_tuning_status)
# Generate OSINT response when prompt is entered
if st.button("Generate OSINT Results"):
if prompt:
response = generate_osint_results(prompt, st.session_state.history)
st.session_state.history.append({"user": prompt, "assistant": response[0]})
st.write("### Generated OSINT Result:")
st.write(response[0])
else:
st.error("Please enter a prompt.")
# Optionally save fine-tuned model
if os.path.exists("./fine_tuned_model"):
st.write("The model has been fine-tuned and saved as `fine_tuned_model`.")
|