Update app.py

app.py

@@ -1,129 +1,206 @@
 import os
-import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
 import time
-
-# Hugging Face API Token
+import re
+import requests
+import json
+from bs4 import BeautifulSoup
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaForCausalLM, LlamaTokenizer, LlamaConfig
+from streamlit_chat import message
+
+# Set page title and icon
+st.set_page_config(page_title="LLaMA Chatbot", page_icon=":robot_face:")
+
+# Custom CSS for styling
+st.markdown(
+    """
+    <style>
+    .stButton>button {
+        background-color: #4CAF50;
+        color: white;
+        border-radius: 12px;
+        padding: 10px 24px;
+    }
+    .stTextArea textarea {
+        background-color: #f5f5f5;
+    }
+    .stDownloadButton>button {
+        background-color: #4CAF50;
+        color: white;
+    }
+    </style>
+    """, unsafe_allow_html=True
+)
+
+# Load Hugging Face API token
 hf_token = os.getenv("HUGGING_FACE_API_TOKEN")
-
-if hf_token:
-    st.write(f"Hugging Face API token found: {hf_token[:4]}...")  # Displaying only the first 4 characters for security
-else:
+if not hf_token:
     st.error("Hugging Face API token not found. Please set the HUGGING_FACE_API_TOKEN environment variable.")
     st.stop()
 
-# Model ID (use a valid model from Hugging Face)
-model_id = "gpt2"  # Replace with a valid model
-
-# Initialize the model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=hf_token)
-model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=hf_token)
-
-# Set pad_token_id to eos_token_id to avoid the warning
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-
-# Alternatively, add a new padding token if it's not defined
-# if tokenizer.pad_token is None:
-#     tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-#     model.resize_token_embeddings(len(tokenizer))
-
-def generate_response(prompt):
-    # Tokenize the prompt with attention mask
-    inputs = tokenizer(prompt, return_tensors='pt', padding=True, truncation=True)
+# Initialize session state variables
+if 'generated' not in st.session_state:
+    st.session_state['generated'] = []
+if 'past' not in st.session_state:
+    st.session_state['past'] = []
+if 'messages' not in st.session_state:
+    st.session_state['messages'] = [{"role": "system", "content": "You are a helpful assistant."}]
+if 'model_name' not in st.session_state:
+    st.session_state['model_name'] = []
+if 'total_tokens' not in st.session_state:
+    st.session_state['total_tokens'] = []
+if 'total_cost' not in st.session_state:
+    st.session_state['total_cost'] = 0.0
+if 'chat_data' not in st.session_state:
+    st.session_state['chat_data'] = []  # For storing the chat logs
+
+# Sidebar - Model Selection, Style Parameters, and Cost Display
+st.sidebar.title("Model Selection")
+model_name = st.sidebar.selectbox("Choose a model:", ["gpt2", "gpt-neo-125M", "distilgpt2", "LLaMA"])
+
+# Parameters to adjust the response style and creativity
+st.sidebar.title("Response Style Controls")
+temperature = st.sidebar.slider("Creativity (Temperature)", min_value=0.0, max_value=1.5, value=0.7, step=0.1)
+top_p = st.sidebar.slider("Nucleus Sampling (Top-p)", min_value=0.5, max_value=1.0, value=0.9, step=0.05)
+top_k = st.sidebar.slider("Token Sampling (Top-k)", min_value=1, max_value=100, value=50, step=1)
+repetition_penalty = st.sidebar.slider("Repetition Penalty", min_value=1.0, max_value=2.0, value=1.2, step=0.1)
+max_length = st.sidebar.slider("Max Length", min_value=50, max_value=4859, value=500, step=10)
+
+# Function to load the model and tokenizer
+@st.cache_resource
+def load_model_and_tokenizer(model_name):
+    if "LLaMA" in model_name:
+        tokenizer = LlamaTokenizer.from_pretrained(model_name, token=hf_token)
+        config = LlamaConfig.from_pretrained(model_name, token=hf_token)
+        model = LlamaForCausalLM.from_pretrained(model_name, config=config, token=hf_token)
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
+        model = AutoModelForCausalLM.from_pretrained(model_name, token=hf_token)
     
-    # Generate text with the attention mask
-    output = model.generate(
-        inputs['input_ids'], 
-        attention_mask=inputs['attention_mask'],  # Pass attention mask to prevent the warning
-        max_length=150, 
-        num_return_sequences=1, 
-        do_sample=True, 
-        top_k=50, 
-        top_p=0.95
-    )
+    return tokenizer, model
+
+tokenizer, model = load_model_and_tokenizer(model_name)
+
+# Function to reset the session
+def reset_session():
+    st.session_state['generated'] = []
+    st.session_state['past'] = []
+    st.session_state['messages'] = [{"role": "system", "content": "You are a helpful assistant."}]
+    st.session_state['model_name'] = []
+    st.session_state['total_tokens'] = []
+    st.session_state['total_cost'] = 0.0
+    st.session_state['chat_data'] = []  # Reset chat logs
+
+# Reset chat button in sidebar
+reset_button = st.sidebar.button("Reset Chat")
+if reset_button:
+    reset_session()
+
+# Function to fetch and parse a webpage for specific tags
+def fetch_website_content(url):
+    try:
+        response = requests.get(url)
+        if response.status_code == 200:
+            soup = BeautifulSoup(response.text, 'html.parser')
+            headings = [h.get_text() for h in soup.find_all(['h1', 'h2', 'h3'])]
+            paragraphs = [p.get_text() for p in soup.find_all('p')]
+            articles = [article.get_text() for article in soup.find_all('article')]
+
+            content = {
+                "headings": headings,
+                "paragraphs": paragraphs,
+                "articles": articles
+            }
+            return content
+        else:
+            return {"error": f"Failed to retrieve content, status code: {response.status_code}"}
+    except Exception as e:
+        return {"error": f"An error occurred: {str(e)}"}
+
+# Function to check if the input contains a URL
+def extract_url_from_text(text):
+    url_pattern = r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+'
+    urls = re.findall(url_pattern, text)
+    return urls
+
+# Function to generate a response using the model with adjustable parameters
+def generate_response(prompt):
+    urls = extract_url_from_text(prompt)
     
-    # Decode the generated output
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    return response
-
-def save_chat():
-    chat_dir = './Intermediate-Chats'
-    if not os.path.exists(chat_dir):
-        os.makedirs(chat_dir)
-    if st.session_state['messages']:
-        filename = f'{chat_dir}/chat_{int(time.time())}.txt'
-        with open(filename, 'w') as f:
-            for message in st.session_state['messages']:
-                f.write(f"{message['role']}: {message['content']}\n")
-        st.session_state['messages'].clear()
-        st.success("Chat saved successfully.")
+    if urls:
+        # If a URL is detected, crawl the webpage and extract content
+        url_content = fetch_website_content(urls[0])  # Crawl only the first URL for simplicity
+        if 'error' in url_content:
+            return url_content['error']
+        else:
+            return f"Headings: {url_content['headings']}\n\nParagraphs: {url_content['paragraphs']}\n\nArticles: {url_content['articles']}"
     else:
-        st.warning("No chat messages to save.")
-
-def load_saved_chats():
-    chat_dir = './Intermediate-Chats'
-    if os.path.exists(chat_dir):
-        files = os.listdir(chat_dir)
-        files.sort(key=lambda x: os.path.getmtime(os.path.join(chat_dir, x)), reverse=True)
-        for file_name in files:
-            display_name = file_name[:-4] if file_name.endswith('.txt') else file_name
-            if st.sidebar.button(display_name):
-                load_chat(os.path.join(chat_dir, file_name))
-
-def load_chat(file_path):
-    st.session_state['messages'].clear()
-    with open(file_path, 'r') as file:
-        for line in file:
-            if ': ' in line:
-                role, content = line.strip().split(': ', 1)
-                st.session_state['messages'].append({'role': role, 'content': content})
-
-def response_generator(content):
-    current_output = ""
-    for word in content.split():
-        current_output += word + " "
-        yield current_output.strip()
-        time.sleep(0.2)
-
-def main():
-    st.title("LLaMA Chat Interface")
-
-    if 'messages' not in st.session_state:
-        st.session_state['messages'] = []
-
-    # Display chat messages
-    for msg in st.session_state.messages:
-        role = msg['role']
-        with st.chat_message(role):
-            st.write(msg['content'])
-
-    # Accept user input
-    user_input = st.chat_input("Enter your prompt:")
-    if user_input:
-        st.session_state.messages.append({"role": "user", "content": user_input})
-        response = generate_response(user_input)
-        st.session_state.messages.append({"role": "assistant", "content": response})
-
-        # Streaming response in the chat interface
-        with st.chat_message("assistant"):
-            placeholder = st.empty()
-            full_response = ""
-            for word in response_generator(response):
-                full_response += word
-                placeholder.write(full_response)
-
-    # Sidebar functionality
-    if st.sidebar.button("Save Chat"):
-        save_chat()
-
-    if st.sidebar.button("New Chat"):
-        st.session_state['messages'].clear()
-
-    if st.sidebar.checkbox("Show/hide chat history"):
-        st.sidebar.title("Previous Chats")
-        load_saved_chats()
-
-if __name__ == "__main__":
-    main()
+        # If no URL, proceed with generating a response from the model
+        inputs = tokenizer(prompt, return_tensors="pt")
+        
+        # Pass attention_mask and set pad_token_id
+        outputs = model.generate(
+            inputs.input_ids,
+            attention_mask=inputs.attention_mask,
+            max_length=max_length,
+            do_sample=True,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            repetition_penalty=repetition_penalty,
+            pad_token_id=tokenizer.eos_token_id  # Set pad_token_id
+        )
+        
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response
+
+# Function to save chat logs for later fine-tuning
+def save_chat_data(chat_data):
+    with open('chat_data.json', 'w') as f:
+        json.dump(chat_data, f, indent=4)
+
+# Containers for chat history and user input
+response_container = st.container()
+container = st.container()
+
+with container:
+    with st.form(key='user_input_form'):
+        user_input = st.text_area("You:", key='user_input', height=100)
+        submit_button = st.form_submit_button("Send")
+
+    if submit_button and user_input:
+        start_time = time.time()
+        output = generate_response(user_input)
+        end_time = time.time()
+        inference_time = end_time - start_time
+
+        # Append user input and model output to session state
+        st.session_state['past'].append(user_input)
+        st.session_state['generated'].append(output)
+        st.session_state['model_name'].append(model_name)
+
+        # Log chat data for future training
+        st.session_state['chat_data'].append({
+            "user_input": user_input,
+            "model_response": output
+        })
+
+        # Save chat data to a file (this could be used later for training)
+        save_chat_data(st.session_state['chat_data'])
+
+        # Calculate tokens and cost
+        num_tokens = len(tokenizer.encode(user_input)) + len(tokenizer.encode(output))
+        st.session_state['total_tokens'].append(num_tokens)
+        cost_per_1000_tokens = 0.0001
+        cost = cost_per_1000_tokens * (num_tokens / 1000)
+        st.session_state['total_cost'] += cost
+
+        # Display chat history
+        with response_container:
+            for i in range(len(st.session_state['generated'])):
+                message(st.session_state['past'][i], is_user=True, key=str(i) + '_user')
+                message(st.session_state['generated'][i], key=str(i))
+                st.write(f"Model: {st.session_state['model_name'][i]}")
+                st.write(f"Tokens: {st.session_state['total_tokens'][i]}")
+                st.write(f"Inference Time: {inference_time:.4f} seconds")
+                st.write(f"Cost: ${cost:.5f}")