Update app.py

app.py

@@ -1,11 +1,175 @@
+import os
+import re
+import requests
 import gradio as gr
-from huggingface_hub import InferenceClient
+from bs4 import BeautifulSoup
+import torch
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+# 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,
@@ -13,52 +177,73 @@ def respond(
     system_message,
     max_tokens,
     temperature,
-    top_p,
+    top_p
 ):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
+    """
+    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
 
-    response = ""
+    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: "
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
+    # Query the local CIAgent
+    response_text = agent.query(
+        prompt,
+        max_length=max_tokens,
         temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+        top_p=top_p
+    )
 
-        response += token
-        yield response
+    # 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
 
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
+# Create the ChatInterface
 demo = gr.ChatInterface(
-    respond,
+    fn=respond,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        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.95,
+            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()