import os import re import requests import gradio as gr from bs4 import BeautifulSoup import torch # Hugging Face Transformers from transformers import ( AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, pipeline ) from datasets import Dataset # ----------------------------- # 1) SCRAPING (OPTIONAL) # ----------------------------- BASE_URL = "https://www.cia.gov" ARCHIVE_URL = "https://www.cia.gov/resources/csi/studies-in-intelligence/archives/operations-subject-index/" def get_article_links(): """ Fetch the archive page and extract article links pointing to the CIA Studies in Intelligence. """ response = requests.get(ARCHIVE_URL) response.raise_for_status() soup = BeautifulSoup(response.text, 'html.parser') links = [] for a_tag in soup.find_all('a', href=True): href = a_tag['href'] if "resources/csi/studies-in-intelligence" in href.lower(): # Convert relative links to absolute if href.startswith("/"): href = BASE_URL + href links.append(href) return list(set(links)) # remove duplicates def scrape_article_text(url): """ Fetch the article text from the URL if it's HTML. (Skipping PDFs for demo.) """ response = requests.get(url) response.raise_for_status() content_type = response.headers.get('Content-Type', '') if 'application/pdf' in content_type.lower(): # Skip PDFs in this simple demo. return None soup = BeautifulSoup(response.text, 'html.parser') paragraphs = soup.find_all('p') article_text = "\n".join(p.get_text(strip=True) for p in paragraphs) return article_text def scrape_all_articles(article_links): """ Iterate through all links and gather text into a dict {url: text}. """ corpus_data = {} for link in article_links: text = scrape_article_text(link) if text: corpus_data[link] = text return corpus_data # ----------------------------- # 2) DATA CLEANING # ----------------------------- import re def clean_text(text): # Simple cleaning: remove extra whitespace text = re.sub(r'\s+', ' ', text) text = text.strip() return text def prepare_dataset(corpus_data): cleaned_texts = [] for url, text in corpus_data.items(): cleaned_texts.append(clean_text(text)) return cleaned_texts # ----------------------------- # 3) FINE-TUNING (OPTIONAL) # ----------------------------- def fine_tune_model(cleaned_texts, model_name="gpt2", output_dir="cia_agent_model"): """ Fine-tunes GPT-2 on your CIA corpus. Warning: resource-intensive! The free Hugging Face Spaces might time out. """ ds = Dataset.from_dict({"text": cleaned_texts}) tokenizer = AutoTokenizer.from_pretrained(model_name) tokenizer.pad_token = tokenizer.eos_token # GPT-2 doesn't have a pad token def tokenize_function(examples): return tokenizer( examples["text"], truncation=True, padding="max_length", max_length=128 ) tokenized_ds = ds.map(tokenize_function, batched=True) model = AutoModelForCausalLM.from_pretrained(model_name) training_args = TrainingArguments( output_dir=output_dir, overwrite_output_dir=True, num_train_epochs=1, # demonstration only per_device_train_batch_size=1, save_steps=100, save_total_limit=1, logging_steps=10, evaluation_strategy="no", # or 'steps' ) trainer = Trainer( model=model, args=training_args, train_dataset=tokenized_ds, ) trainer.train() trainer.save_model(output_dir) tokenizer.save_pretrained(output_dir) return model, tokenizer # ----------------------------- # 4) CIAgent INFERENCE # ----------------------------- class CIAgent: def __init__(self, model_path="cia_agent_model"): """ Initialize a pipeline from a local fine-tuned model folder or fallback to GPT-2. """ if not os.path.exists(model_path): model_path = "gpt2" self.generator = pipeline( "text-generation", model=model_path, tokenizer=model_path ) self.max_length = 128 def query(self, prompt, max_length=128, temperature=0.7, top_p=0.9): """ Generate text from the model. """ outputs = self.generator( prompt, max_length=max_length, temperature=temperature, top_p=top_p, num_return_sequences=1 ) return outputs[0]["generated_text"] # ----------------------------- # 5) GRADIO CHAT INTERFACE # ----------------------------- # Create (or load) your CIAgent. In a real workflow, you might have already # fine-tuned locally and just upload the "cia_agent_model" folder to your Space. agent = CIAgent(model_path="cia_agent_model") # or "gpt2" if you haven't trained def respond( message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p ): """ This function is called by Gradio's ChatInterface. It receives: - message: current user message - history: list of (user_text, assistant_text) pairs - system_message: the "system" instruction to guide the model - max_tokens, temperature, top_p: generation parameters We build a 'prompt' from all conversation turns + system message. Then we query the CIAgent to get one text output. Since CIAgent doesn't stream tokens by default, we yield once at the end. """ # Build the conversation prompt # For demonstration, we simply concatenate everything in a naive format. # You could style it in a more advanced way for better context handling. prompt = f"System: {system_message}\n\n" for user_text, assistant_text in history: if user_text: prompt += f"User: {user_text}\n" if assistant_text: prompt += f"Assistant: {assistant_text}\n" # Add the new user message prompt += f"User: {message}\nAssistant: " # Query the local CIAgent response_text = agent.query( prompt, max_length=max_tokens, temperature=temperature, top_p=top_p ) # We can yield partial tokens if we want streaming, but the pipeline # returns the entire text at once. Let's yield a single chunk: yield response_text # Create the ChatInterface demo = gr.ChatInterface( fn=respond, additional_inputs=[ gr.Textbox( value="You are a friendly Chatbot that knows about CIA Studies in Intelligence.", label="System message" ), gr.Slider(minimum=1, maximum=2048, value=256, step=1, label="Max new tokens"), gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"), gr.Slider( minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p (nucleus sampling)", ), ], ) if __name__ == "__main__": # WARNING: Running scraping & fine-tuning on a free HF Space # might exceed time/memory limits. If you do want to train, uncomment: # # article_links = get_article_links() # corpus_data = scrape_all_articles(article_links) # cleaned_texts = prepare_dataset(corpus_data) # model, tokenizer = fine_tune_model(cleaned_texts) # # Then re-initialize agent = CIAgent("cia_agent_model") # # For now, just launch the Gradio chat using the existing or fallback GPT-2 model. demo.launch()