Update app.py

app.py

@@ -3,6 +3,8 @@ from bs4 import BeautifulSoup
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 import torch
+import re
+import os
 
 # Web scraping
 def scrape_website(url):
@@ -11,10 +13,15 @@ def scrape_website(url):
     content = ' '.join([p.text for p in soup.find_all('p')])
     return content
 
-# Store data
+# Store data and save to a file
 stored_data = {}
 def store_data(url, content):
     stored_data[url] = content
+    # Save content to a file
+    filename = url.replace("https://", "").replace("http://", "").replace("/", "_") + ".txt"
+    with open(filename, "w") as file:
+        file.write(content)
+    return filename
 
 # Conversational AI with a smaller model
 tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
@@ -24,6 +31,12 @@ model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 model.to(device)
 
+# Function to extract URLs from text
+def extract_urls(text):
+    url_pattern = re.compile(r'(https?://\S+)')
+    urls = url_pattern.findall(text)
+    return urls
+
 def generate_response(input_text):
     input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt').to(device)
     response_ids = model.generate(input_ids, max_length=50, pad_token_id=tokenizer.eos_token_id)
@@ -31,17 +44,28 @@ def generate_response(input_text):
     return response
 
 def chatbot_response(user_input):
-    if user_input.startswith('http'):
-        url = user_input
+    # Extract URLs from the input
+    urls = extract_urls(user_input)
+    
+    file_links = []
+    
+    # Scrape content from URLs, save to file, and append content to user input
+    for url in urls:
         if url in stored_data:
             content = stored_data[url]
+            filename = store_data(url, content)
         else:
             content = scrape_website(url)
-            store_data(url, content)
-        return "I've fetched the data from the website. How can I help you with it?"
-    else:
-        response = generate_response(user_input)
-        return response
+            filename = store_data(url, content)
+        file_links.append(f"You can download the data here: {filename}")
+        user_input += " " + content
+    
+    # Generate response based on the combined input
+    response = generate_response(user_input)
+    
+    # Combine the chatbot response with file links
+    full_response = response + "\n" + "\n".join(file_links)
+    return full_response
 
 # Interface
 def chat_interface(user_input):