|
--- |
|
license: mit |
|
language: |
|
- en |
|
base_model: |
|
- meta-llama/Llama-3.1-8B |
|
pipeline_tag: reinforcement-learning |
|
--- |
|
|
|
import os |
|
import tkinter as tk |
|
from tkinter import filedialog, messagebox |
|
import PyPDF2 |
|
import re |
|
import json |
|
import torch |
|
import ollama |
|
from openai import OpenAI |
|
import argparse |
|
|
|
# ANSI escape codes for colors |
|
PINK = '\033[95m' |
|
CYAN = '\033[96m' |
|
YELLOW = '\033[93m' |
|
NEON_GREEN = '\033[92m' |
|
RESET_COLOR = '\033[0m' |
|
|
|
# Function to open a file and return its contents as a string |
|
def open_file(filepath): |
|
with open(filepath, 'r', encoding='utf-8') as infile: |
|
return infile.read() |
|
|
|
# Function to convert PDF to text and append to vault.txt |
|
def convert_pdf_to_text(): |
|
file_path = filedialog.askopenfilename(filetypes=[("PDF Files", "*.pdf")]) |
|
if file_path: |
|
base_directory = os.path.join("local-rag", "text_parse") |
|
file_name = os.path.basename(file_path) |
|
output_file_name = os.path.splitext(file_name)[0] + ".txt" |
|
file_output_path = os.path.join(base_directory, output_file_name) |
|
|
|
if not os.path.exists(base_directory): |
|
os.makedirs(base_directory) |
|
print(f"Directory '{base_directory}' created.") |
|
|
|
with open(file_path, 'rb') as pdf_file: |
|
pdf_reader = PyPDF2.PdfReader(pdf_file) |
|
text = '' |
|
for page_num in range(len(pdf_reader.pages)): |
|
page = pdf_reader.pages[page_num] |
|
if page.extract_text(): |
|
text += page.extract_text() + " " |
|
|
|
text = re.sub(r'\s+', ' ', text).strip() |
|
sentences = re.split(r'(?<=[.!?]) +', text) |
|
chunks = [] |
|
current_chunk = "" |
|
for sentence in sentences: |
|
if len(current_chunk) + len(sentence) + 1 < 1000: |
|
current_chunk += (sentence + " ").strip() |
|
else: |
|
chunks.append(current_chunk) |
|
current_chunk = sentence + " " |
|
if current_chunk: |
|
chunks.append(current_chunk) |
|
|
|
with open(os.path.join("local-rag", "temp.txt"), "w", encoding="utf-8") as temp_file: |
|
temp_file.write(output_file_name + "\n") |
|
for chunk in chunks: |
|
temp_file.write(chunk.strip() + "\n") |
|
|
|
with open(os.path.join("local-rag", "vault.txt"), "a", encoding="utf-8") as vault_file: |
|
vault_file.write("\n") |
|
for chunk in chunks: |
|
vault_file.write(chunk.strip() + "\n") |
|
|
|
if not os.path.exists(file_output_path): |
|
with open(file_output_path, "w", encoding="utf-8") as f: |
|
for chunk in chunks: |
|
f.write(chunk.strip() + "\n") |
|
f.write("====================NOT FINISHED====================\n") |
|
print(f"File '{file_output_path}' created with NOT FINISHED flag at the end.") |
|
else: |
|
print(f"File '{file_output_path}' already exists.") |
|
|
|
print(f"PDF content appended to vault.txt with each chunk on a separate line.") |
|
# Call the second part after the PDF conversion is done |
|
|
|
input_value = input("Enter your question:") |
|
process_text_files(input_value) |
|
|
|
# Function to upload a text file and append to vault.txt |
|
def upload_txtfile(): |
|
file_path = filedialog.askopenfilename(filetypes=[("Text Files", "*.txt")]) |
|
if file_path: |
|
# Define the base directory |
|
base_directory = os.path.join("local-rag", "text_parse") |
|
|
|
# Get the file name without the directory and extension |
|
file_name = os.path.basename(file_path) |
|
output_file_name = os.path.splitext(file_name)[0] + ".txt" # Convert PDF filename to .txt |
|
|
|
|
|
# Construct the output file path in the base directory |
|
file_output_path = os.path.join(base_directory, output_file_name) |
|
|
|
# Create base directory if it doesn't exist |
|
if not os.path.exists(base_directory): |
|
os.makedirs(base_directory) |
|
print(f"Directory '{base_directory}' created.") |
|
|
|
|
|
with open(file_path, 'r', encoding="utf-8") as txt_file: |
|
text = txt_file.read() |
|
|
|
# Normalize whitespace and clean up text |
|
text = re.sub(r'\s+', ' ', text).strip() |
|
|
|
# Split text into chunks by sentences, respecting a maximum chunk size |
|
sentences = re.split(r'(?<=[.!?]) +', text) # split on spaces following sentence-ending punctuation |
|
chunks = [] |
|
current_chunk = "" |
|
for sentence in sentences: |
|
# Check if the current sentence plus the current chunk exceeds the limit |
|
if len(current_chunk) + len(sentence) + 1 < 1000: # +1 for the space |
|
current_chunk += (sentence + " ").strip() |
|
else: |
|
# When the chunk exceeds 1000 characters, store it and start a new one |
|
chunks.append(current_chunk) |
|
current_chunk = sentence + " " |
|
if current_chunk: # Don't forget the last chunk! |
|
chunks.append(current_chunk) |
|
|
|
# Clear temp.txt and write the new content |
|
with open(os.path.join("local-rag", "temp.txt"), "w", encoding="utf-8") as temp_file: |
|
temp_file.write(output_file_name + "\n") # Write the output file name as the first line |
|
for chunk in chunks: |
|
# Write each chunk to its own line |
|
temp_file.write(chunk.strip() + "\n") # Each chunk on a new line |
|
|
|
with open(os.path.join("local-rag", "vault.txt"), "a", encoding="utf-8") as vault_file: |
|
vault_file.write("\n") # Add a new line to separate content |
|
for chunk in chunks: |
|
# Write each chunk to its own line |
|
vault_file.write(chunk.strip() + "\n") # Two newlines to separate chunks |
|
|
|
# Create the file in the directory if it doesn't exist |
|
if not os.path.exists(file_output_path): |
|
with open(file_output_path, "w") as f: |
|
f.write("") # Create an empty file |
|
f.write("====================NOT FINISHED====================\n") |
|
print(f"File '{file_output_path}' created with NOT FINISHED flag at the end.") |
|
else: |
|
print(f"File '{file_output_path}' already exists.") |
|
|
|
print(f"Text file content appended to vault.txt with each chunk on a separate line.") |
|
|
|
input_value = input("Enter your question:") |
|
process_text_files(input_value) |
|
else: |
|
print("No file selected.") |
|
|
|
# Function to upload a JSON file and append to vault.txt |
|
def upload_jsonfile(): |
|
file_path = filedialog.askopenfilename(filetypes=[("JSON Files", "*.json")]) |
|
if file_path: |
|
|
|
# Define the base directory |
|
base_directory = os.path.join("local-rag", "text_parse") |
|
|
|
# Get the file name without the directory and extension |
|
file_name = os.path.basename(file_path) |
|
output_file_name = os.path.splitext(file_name)[0] + ".txt" # Convert PDF filename to .txt |
|
|
|
|
|
# Construct the output file path in the base directory |
|
file_output_path = os.path.join(base_directory, output_file_name) |
|
|
|
# Create base directory if it doesn't exist |
|
if not os.path.exists(base_directory): |
|
os.makedirs(base_directory) |
|
print(f"Directory '{base_directory}' created.") |
|
|
|
|
|
|
|
|
|
with open(file_path, 'r', encoding="utf-8") as json_file: |
|
data = json.load(json_file) |
|
|
|
# Flatten the JSON data into a single string |
|
text = json.dumps(data, ensure_ascii=False) |
|
|
|
# Normalize whitespace and clean up text |
|
text = re.sub(r'\s+', ' ', text).strip() |
|
|
|
# Split text into chunks by sentences, respecting a maximum chunk size |
|
sentences = re.split(r'(?<=[.!?]) +', text) # split on spaces following sentence-ending punctuation |
|
chunks = [] |
|
current_chunk = "" |
|
for sentence in sentences: |
|
# Check if the current sentence plus the current chunk exceeds the limit |
|
if len(current_chunk) + len(sentence) + 1 < 1000: # +1 for the space |
|
current_chunk += (sentence + " ").strip() |
|
else: |
|
# When the chunk exceeds 1000 characters, store it and start a new one |
|
chunks.append(current_chunk) |
|
current_chunk = sentence + " " |
|
if current_chunk: # Don't forget the last chunk! |
|
chunks.append(current_chunk) |
|
|
|
# Clear temp.txt and write the new content |
|
with open(os.path.join("local-rag", "temp.txt"), "w", encoding="utf-8") as temp_file: |
|
temp_file.write(output_file_name + "\n") # Write the output file name as the first line |
|
for chunk in chunks: |
|
# Write each chunk to its own line |
|
temp_file.write(chunk.strip() + "\n") # Each chunk on a new line |
|
|
|
with open(os.path.join("local-rag", "vault.txt"), "a", encoding="utf-8") as vault_file: |
|
vault_file.write("\n") # Add a new line to separate content |
|
for chunk in chunks: |
|
# Write each chunk to its own line |
|
vault_file.write(chunk.strip() + "\n") # Two newlines to separate chunks |
|
|
|
if not os.path.exists(file_output_path): |
|
with open(file_output_path, "w", encoding="utf-8") as f: |
|
for chunk in chunks: |
|
f.write(chunk.strip() + "\n") # Each chunk on a new line |
|
f.write("====================NOT FINISHED====================\n") |
|
print(f"File '{file_output_path}' created with NOT FINISHED flag at the end.") |
|
else: |
|
print(f"File '{file_output_path}' already exists.") |
|
|
|
|
|
|
|
print(f"JSON file content appended to vault.txt with each chunk on a separate line.") |
|
|
|
input_value = input("Enter your question:") |
|
process_text_files(input_value) |
|
|
|
def summarize(): |
|
summary_window = tk.Toplevel(root) |
|
summary_window.title("Text Summarizer") |
|
summary_window.geometry("400x200") |
|
|
|
# Create a label for the window |
|
label = tk.Label(summary_window, text="Choose an option to summarize text:") |
|
label.pack(pady=10) |
|
|
|
# Create two buttons: one for uploading a .txt file, and one for pasting text directly |
|
upload_button = tk.Button(summary_window, text="Upload from .txt File", command=summarize_from_file) |
|
upload_button.pack(pady=5) |
|
|
|
paste_button = tk.Button(summary_window, text="Paste your text", command=lambda: open_paste_window(summary_window)) |
|
paste_button.pack(pady=5) |
|
|
|
# Function to upload a .txt file and summarize |
|
def summarize_from_file(): |
|
file_path = filedialog.askopenfilename(filetypes=[("Text Files", "*.txt")]) |
|
if file_path: |
|
# Define the base directory where the file will be saved |
|
base_directory = os.path.join("local-rag", "text_sum") |
|
|
|
file_name = os.path.basename(file_path) |
|
|
|
# Create the directory if it doesn't exist |
|
if not os.path.exists(base_directory): |
|
os.makedirs(base_directory) |
|
print(f"Directory '{base_directory}' created.") |
|
|
|
summary_content = [] |
|
if os.path.exists(file_name): |
|
with open(file_name, "r", encoding='utf-8') as sum_file: |
|
summary_content = sum_file.readlines() |
|
|
|
summary_embeddings = [] |
|
for content in summary_content: |
|
response = ollama.embeddings(model='mxbai-embed-large', prompt=content) |
|
summary_embeddings.append(response["embedding"]) |
|
|
|
summary_embeddings_tensor = torch.tensor(summary_embeddings) |
|
print("Embeddings for each line in the vault:") |
|
print(summary_embeddings_tensor) |
|
|
|
conversation_history = [] |
|
system_message = "You are a helpful assistant that is an expert at summarizing the text from a given document" |
|
user_input = "Summarize this paragraph" |
|
|
|
response = ollama_chat(user_input, system_message, summary_embeddings_tensor, summary_content, args.model, conversation_history) |
|
|
|
messagebox.showinfo("Summary", response) # Replace with actual summarizing logic |
|
else: |
|
messagebox.showerror("Error", "No file selected!") |
|
|
|
# Function to open a window for pasting text and summarizing |
|
def open_paste_window(parent_window): |
|
# Create a new window for pasting text |
|
paste_window = tk.Toplevel(parent_window) |
|
paste_window.title("Paste Your Text") |
|
paste_window.geometry("400x300") |
|
|
|
# Create a label and text box for the pasted text |
|
label = tk.Label(paste_window, text="Paste your text below:") |
|
label.pack(pady=5) |
|
|
|
input_textbox = tk.Text(paste_window, height=8, width=40) |
|
input_textbox.pack(pady=5) |
|
|
|
# Function to handle the "Submit" button click |
|
def submit_text(): |
|
pasted_text = input_textbox.get("1.0", tk.END).strip() |
|
if pasted_text: |
|
|
|
system_message = "You are a helpful assistant that is an expert at summarizing the text from a given document" |
|
user_input = "Summarize this paragraph:" |
|
new_value = user_input + pasted_text |
|
messages = [ |
|
{ |
|
"system", |
|
system_message, |
|
}, |
|
{"human", new_value}, |
|
] |
|
response = client.chat.completions.create(model=args.model, messages=messages) |
|
|
|
response_value = response.choices[0].message.content |
|
|
|
|
|
messagebox.showinfo("Summary", response_value) # Replace with actual summarizing logic |
|
paste_window.destroy() # Close the window |
|
else: |
|
messagebox.showerror("Error", "No text entered!") |
|
|
|
# Add Submit and Cancel buttons |
|
submit_button = tk.Button(paste_window, text="Submit", command=submit_text) |
|
submit_button.pack(side=tk.LEFT, padx=10, pady=10) |
|
|
|
cancel_button = tk.Button(paste_window, text="Cancel", command=paste_window.destroy) |
|
cancel_button.pack(side=tk.RIGHT, padx=10, pady=10) |
|
|
|
|
|
# Function to get relevant context from the vault based on user input |
|
def get_relevant_context(rewritten_input, vault_embeddings, vault_content, top_k=3): |
|
if vault_embeddings.nelement() == 0: |
|
return [] |
|
input_embedding = ollama.embeddings(model='mxbai-embed-large', prompt=rewritten_input)["embedding"] |
|
cos_scores = torch.cosine_similarity(torch.tensor(input_embedding).unsqueeze(0), vault_embeddings) |
|
top_k = min(top_k, len(cos_scores)) |
|
top_indices = torch.topk(cos_scores, k=top_k)[1].tolist() |
|
relevant_context = [vault_content[idx].strip() for idx in top_indices] |
|
return relevant_context |
|
|
|
# Function to interact with the Ollama model |
|
def ollama_chat(user_input, system_message, vault_embeddings, vault_content, ollama_model, conversation_history): |
|
relevant_context = get_relevant_context(user_input, vault_embeddings, vault_content, top_k=3) |
|
if relevant_context: |
|
context_str = "\n".join(relevant_context) |
|
print("Context Pulled from Documents: \n\n" + CYAN + context_str + RESET_COLOR) |
|
else: |
|
print(CYAN + "No relevant context found." + RESET_COLOR) |
|
|
|
user_input_with_context = user_input |
|
if relevant_context: |
|
user_input_with_context = context_str + "\n\n" + user_input |
|
|
|
conversation_history.append({"role": "user", "content": user_input_with_context}) |
|
messages = [{"role": "system", "content": system_message}, *conversation_history] |
|
|
|
response = client.chat.completions.create(model=ollama_model, messages=messages) |
|
conversation_history.append({"role": "assistant", "content": response.choices[0].message.content}) |
|
|
|
return response.choices[0].message.content |
|
|
|
# Function to process text files, check for NOT FINISHED flag, and compute embeddings |
|
def process_text_files(user_input): |
|
text_parse_directory = os.path.join("local-rag", "text_parse") |
|
temp_file_path = os.path.join("local-rag", "temp.txt") |
|
|
|
if not os.path.exists(text_parse_directory): |
|
print(f"Directory '{text_parse_directory}' does not exist.") |
|
return False |
|
|
|
if not os.path.exists(temp_file_path): |
|
print("temp.txt does not exist.") |
|
return False |
|
|
|
with open(temp_file_path, 'r', encoding='utf-8') as temp_file: |
|
first_line = temp_file.readline().strip() |
|
|
|
text_files = [f for f in os.listdir(text_parse_directory) if f.endswith('.txt')] |
|
|
|
if f"{first_line}" not in text_files: |
|
print(f"No matching file found for '{first_line}.txt' in text_parse directory.") |
|
return False |
|
|
|
file_path = os.path.join(text_parse_directory, f"{first_line}") |
|
with open(file_path, 'r', encoding='utf-8') as f: |
|
lines = f.readlines() |
|
|
|
lines = [line.strip() for line in lines] |
|
|
|
if len(lines) >= 2 and lines[-1] == "====================NOT FINISHED====================": |
|
print(f"'{first_line}' contains the 'NOT FINISHED' flag. Computing embeddings.") |
|
|
|
vault_content = [] |
|
if os.path.exists(temp_file_path): |
|
with open(temp_file_path, "r", encoding='utf-8') as vault_file: |
|
vault_content = vault_file.readlines() |
|
|
|
vault_embeddings = [] |
|
for content in vault_content: |
|
response = ollama.embeddings(model='mxbai-embed-large', prompt=content) |
|
vault_embeddings.append(response["embedding"]) |
|
|
|
vault_embeddings_tensor = torch.tensor(vault_embeddings) |
|
print("Embeddings for each line in the vault:") |
|
print(vault_embeddings_tensor) |
|
|
|
with open(os.path.join(text_parse_directory, f"{first_line}_embedding.pt"), "wb") as tensor_file: |
|
torch.save(vault_embeddings_tensor, tensor_file) |
|
|
|
with open(file_path, 'w', encoding='utf-8') as f: |
|
f.writelines(lines[:-1]) |
|
|
|
else: |
|
print(f"'{first_line}' does not contain the 'NOT FINISHED' flag or is already complete. Loading tensor if it exists.") |
|
|
|
tensor_file_path = os.path.join(text_parse_directory, f"{first_line}_embedding.pt") |
|
if os.path.exists(tensor_file_path): |
|
vault_embeddings_tensor = torch.load(tensor_file_path) |
|
print("Loaded Vault Embedding Tensor:") |
|
print(vault_embeddings_tensor) |
|
|
|
vault_content = [] |
|
if os.path.exists(temp_file_path): |
|
with open(temp_file_path, "r", encoding='utf-8') as vault_file: |
|
vault_content = vault_file.readlines() |
|
|
|
conversation_history = [] |
|
system_message = "You are a helpful assistant that is an expert at extracting the most useful information from a given text" |
|
response = ollama_chat(user_input, system_message, vault_embeddings_tensor, vault_content, args.model, conversation_history) |
|
|
|
print (response) |
|
|
|
return response |
|
|
|
# Create the main window |
|
root = tk.Tk() |
|
root.title("Upload .pdf, .txt, or .json") |
|
|
|
# Create a button to open the file dialog for PDF |
|
pdf_button = tk.Button(root, text="Upload PDF", command=convert_pdf_to_text) |
|
pdf_button.pack(pady=15) |
|
|
|
# Create a button to open the file dialog for text file |
|
txt_button = tk.Button(root, text="Upload Text File", command=upload_txtfile) |
|
txt_button.pack(pady=15) |
|
|
|
# Create a button to open the file dialog for JSON file |
|
json_button = tk.Button(root, text="Upload JSON File", command=upload_jsonfile) |
|
json_button.pack(pady=15) |
|
|
|
# Create a button to open the summerizer |
|
json_button = tk.Button(root, text="Summarize This!", command=summarize) |
|
json_button.pack(pady=15) |
|
|
|
# Configuration for the Ollama API client |
|
client = OpenAI(base_url='http://localhost:11434/v1', api_key='llama3') |
|
|
|
# Parse command-line arguments |
|
parser = argparse.ArgumentParser(description="Ollama Chat") |
|
parser.add_argument("--model", default="llama3", help="Ollama model to use (default: llama3)") |
|
args = parser.parse_args() |
|
|
|
# Run the main event loop |
|
root.mainloop() |
|
|
