File size: 37,472 Bytes
f4368a7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 |
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
import datetime
import functools
import traceback
from typing import List, Optional, Any, Dict
import torch
import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from langchain_community.llms import HuggingFacePipeline
# Other LangChain and community imports
from langchain_community.document_loaders import OnlinePDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import EnsembleRetriever
from langchain.prompts import ChatPromptTemplate
from langchain.schema import StrOutputParser, Document
from langchain_core.runnables import RunnableParallel, RunnableLambda
from transformers.quantizers.auto import AutoQuantizationConfig
import gradio as gr
import requests
from pydantic import PrivateAttr
import pydantic
from langchain.llms.base import LLM
from typing import Any, Optional, List
import typing
import time
print("Pydantic Version: ")
print(pydantic.__version__)
# Add Mistral imports with fallback handling
try:
from mistralai import Mistral
MISTRAL_AVAILABLE = True
debug_print = lambda msg: print(f"[{datetime.datetime.now().isoformat()}] {msg}")
debug_print("Loaded latest Mistral client library")
except ImportError:
MISTRAL_AVAILABLE = False
debug_print = lambda msg: print(f"[{datetime.datetime.now().isoformat()}] {msg}")
debug_print("Mistral client library not found. Install with: pip install mistralai")
def debug_print(message: str):
print(f"[{datetime.datetime.now().isoformat()}] {message}", flush=True)
def word_count(text: str) -> int:
return len(text.split())
# Initialize a tokenizer for token counting (using gpt2 as a generic fallback)
def initialize_tokenizer():
try:
return AutoTokenizer.from_pretrained("gpt2")
except Exception as e:
debug_print("Failed to initialize tokenizer: " + str(e))
return None
global_tokenizer = initialize_tokenizer()
def count_tokens(text: str) -> int:
if global_tokenizer:
try:
return len(global_tokenizer.encode(text))
except Exception as e:
return len(text.split())
return len(text.split())
# Add these imports at the top of your file
import uuid
import threading
import queue
from typing import Dict, Any, Tuple, Optional
import time
# Global storage for jobs and results
jobs = {} # Stores job status and results
results_queue = queue.Queue() # Thread-safe queue for completed jobs
processing_lock = threading.Lock() # Prevent simultaneous processing of the same job
# Add a global variable to store the last job ID
last_job_id = None
# Add these missing async processing functions
def process_in_background(job_id, function, args):
"""Process a function in the background and store results"""
try:
debug_print(f"Processing job {job_id} in background")
result = function(*args)
results_queue.put((job_id, result))
debug_print(f"Job {job_id} completed and added to results queue")
except Exception as e:
debug_print(f"Error in background job {job_id}: {str(e)}")
error_result = (f"Error processing job: {str(e)}", "", "", "")
results_queue.put((job_id, error_result))
def load_pdfs_async(file_links, model_choice, prompt_template, bm25_weight, temperature, top_p):
"""Asynchronous version of load_pdfs_updated to prevent timeouts"""
global last_job_id
if not file_links:
return "Please enter non-empty URLs", "", "Model used: N/A", "", "", get_job_list()
job_id = str(uuid.uuid4())
debug_print(f"Starting async job {job_id} for file loading")
# Start background thread
threading.Thread(
target=process_in_background,
args=(job_id, load_pdfs_updated, [file_links, model_choice, prompt_template, bm25_weight, temperature, top_p])
).start()
job_query = f"Loading files: {file_links.split()[0]}..." if file_links else "No files"
jobs[job_id] = {
"status": "processing",
"type": "load_files",
"start_time": time.time(),
"query": job_query
}
last_job_id = job_id
return (
f"Files submitted and processing in the background (Job ID: {job_id}).\n\n"
f"Use 'Check Job Status' tab with this ID to get results.",
f"Job ID: {job_id}",
f"Model requested: {model_choice}",
job_id, # Return job_id to update the job_id_input component
job_query, # Return job_query to update the job_query_display component
get_job_list() # Return updated job list
)
def submit_query_async(query, use_llama, use_mistral, temperature, top_p):
"""Asynchronous version of submit_query_updated to prevent timeouts"""
global last_job_id
if not query:
return ("Please enter a non-empty query", "Input/Output tokens: 0/0",
"Please enter a non-empty query", "Input/Output tokens: 0/0",
"", "", get_job_list())
if not (use_llama or use_mistral):
return ("Please select at least one model", "Input/Output tokens: 0/0",
"Please select at least one model", "Input/Output tokens: 0/0",
"", "", get_job_list())
responses = {"llama": None, "mistral": None}
job_ids = []
if use_llama:
llama_job_id = str(uuid.uuid4())
debug_print(f"Starting async job {llama_job_id} for Llama query: {query}")
# Start background thread for Llama
threading.Thread(
target=process_in_background,
args=(llama_job_id, submit_query_updated, [query, "๐บ๐ธ Remote Meta-Llama-3", temperature, top_p])
).start()
jobs[llama_job_id] = {
"status": "processing",
"type": "query",
"start_time": time.time(),
"query": query,
"model": "Llama"
}
job_ids.append(llama_job_id)
responses["llama"] = f"Processing (Job ID: {llama_job_id})"
if use_mistral:
mistral_job_id = str(uuid.uuid4())
debug_print(f"Starting async job {mistral_job_id} for Mistral query: {query}")
# Start background thread for Mistral
threading.Thread(
target=process_in_background,
args=(mistral_job_id, submit_query_updated, [query, "๐ช๐บ Mistral-API", temperature, top_p])
).start()
jobs[mistral_job_id] = {
"status": "processing",
"type": "query",
"start_time": time.time(),
"query": query,
"model": "Mistral"
}
job_ids.append(mistral_job_id)
responses["mistral"] = f"Processing (Job ID: {mistral_job_id})"
# Store the last job ID (use the first one for now)
last_job_id = job_ids[0] if job_ids else None
return (
responses.get("llama", "Not selected"),
"Input tokens: " + str(count_tokens(query)) if use_llama else "Not selected",
responses.get("mistral", "Not selected"),
"Input tokens: " + str(count_tokens(query)) if use_mistral else "Not selected",
last_job_id,
query,
get_job_list()
)
def update_ui_with_last_job_id():
# This function doesn't need to do anything anymore
# We'll update the UI directly in the functions that call this
pass
# Function to display all jobs as a clickable list
def get_job_list():
job_list_md = "### Submitted Jobs\n\n"
if not jobs:
return "No jobs found. Submit a query or load files to create jobs."
# Sort jobs by start time (newest first)
sorted_jobs = sorted(
[(job_id, job_info) for job_id, job_info in jobs.items()],
key=lambda x: x[1].get("start_time", 0),
reverse=True
)
for job_id, job_info in sorted_jobs:
status = job_info.get("status", "unknown")
job_type = job_info.get("type", "unknown")
query = job_info.get("query", "")
model = job_info.get("model", "") # Get the model name
start_time = job_info.get("start_time", 0)
time_str = datetime.datetime.fromtimestamp(start_time).strftime("%Y-%m-%d %H:%M:%S")
# Create a shortened query preview
query_preview = query[:30] + "..." if query and len(query) > 30 else query or "N/A"
# Add color and icons based on status
if status == "processing":
status_formatted = f"<span style='color: red'>โณ {status}</span>"
elif status == "completed":
status_formatted = f"<span style='color: green'>โ
{status}</span>"
else:
status_formatted = f"<span style='color: orange'>โ {status}</span>"
# Add model icon based on model name
model_icon = "๐บ๐ธ" if model == "Llama" else "๐ช๐บ" if model == "Mistral" else ""
model_prefix = f"{model_icon} {model} " if model else ""
# Create clickable links using Markdown
if job_type == "query":
job_list_md += f"- [{job_id}](javascript:void) - {time_str} - {status_formatted} - {model_prefix}Query: {query_preview}\n"
else:
job_list_md += f"- [{job_id}](javascript:void) - {time_str} - {status_formatted} - File Load Job\n"
return job_list_md
# Function to handle job list clicks
def job_selected(job_id):
if job_id in jobs:
return job_id, jobs[job_id].get("query", "No query for this job")
return job_id, "Job not found"
# Function to refresh the job list
def refresh_job_list():
return get_job_list()
# Function to sync model dropdown boxes
def sync_model_dropdown(value):
return value
# Function to check job status
def check_job_status(job_id):
if not job_id:
return "Please enter a job ID", "", "", "", ""
# Process any completed jobs in the queue
try:
while not results_queue.empty():
completed_id, result = results_queue.get_nowait()
if completed_id in jobs:
jobs[completed_id]["status"] = "completed"
jobs[completed_id]["result"] = result
jobs[completed_id]["end_time"] = time.time()
debug_print(f"Job {completed_id} completed and stored in jobs dictionary")
except queue.Empty:
pass
# Check if the requested job exists
if job_id not in jobs:
return "Job not found. Please check the ID and try again.", "", "", "", ""
job = jobs[job_id]
job_query = job.get("query", "No query available for this job")
# If job is still processing
if job["status"] == "processing":
elapsed_time = time.time() - job["start_time"]
job_type = job.get("type", "unknown")
if job_type == "load_files":
return (
f"Files are still being processed (elapsed: {elapsed_time:.1f}s).\n\n"
f"Try checking again in a few seconds.",
f"Job ID: {job_id}",
f"Status: Processing",
"",
job_query
)
else: # query job
return (
f"Query is still being processed (elapsed: {elapsed_time:.1f}s).\n\n"
f"Try checking again in a few seconds.",
f"Job ID: {job_id}",
f"Input tokens: {count_tokens(job.get('query', ''))}",
"Output tokens: pending",
job_query
)
# If job is completed
if job["status"] == "completed":
result = job["result"]
processing_time = job["end_time"] - job["start_time"]
if job.get("type") == "load_files":
return (
f"{result[0]}\n\nProcessing time: {processing_time:.1f}s",
result[1],
result[2],
"",
job_query
)
else: # query job
return (
f"{result[0]}\n\nProcessing time: {processing_time:.1f}s",
result[1],
result[2],
result[3],
job_query
)
# Fallback for unknown status
return f"Job status: {job['status']}", "", "", "", job_query
# Function to clean up old jobs
def cleanup_old_jobs():
current_time = time.time()
to_delete = []
for job_id, job in jobs.items():
# Keep completed jobs for 24 hours, processing jobs for 48 hours
if job["status"] == "completed" and (current_time - job.get("end_time", 0)) > 86400:
to_delete.append(job_id)
elif job["status"] == "processing" and (current_time - job.get("start_time", 0)) > 172800:
to_delete.append(job_id)
for job_id in to_delete:
del jobs[job_id]
debug_print(f"Cleaned up {len(to_delete)} old jobs. {len(jobs)} jobs remaining.")
return f"Cleaned up {len(to_delete)} old jobs", "", ""
# Improve the truncate_prompt function to be more aggressive with limiting context
def truncate_prompt(prompt: str, max_tokens: int = 4096) -> str:
"""Truncate prompt to fit within token limit, preserving the most recent/relevant parts."""
if not prompt:
return ""
if global_tokenizer:
try:
tokens = global_tokenizer.encode(prompt)
if len(tokens) > max_tokens:
# For prompts, we often want to keep the beginning instructions and the end context
# So we'll keep the first 20% and the last 80% of the max tokens
beginning_tokens = int(max_tokens * 0.2)
ending_tokens = max_tokens - beginning_tokens
new_tokens = tokens[:beginning_tokens] + tokens[-(ending_tokens):]
return global_tokenizer.decode(new_tokens)
except Exception as e:
debug_print(f"Truncation error: {str(e)}")
# Fallback to word-based truncation
words = prompt.split()
if len(words) > max_tokens:
beginning_words = int(max_tokens * 0.2)
ending_words = max_tokens - beginning_words
return " ".join(words[:beginning_words] + words[-(ending_words):])
return prompt
default_prompt = """\
{conversation_history}
Use the following context to provide a detailed technical answer to the user's question.
Do not include an introduction like "Based on the provided documents, ...". Just answer the question.
If you don't know the answer, please respond with "I don't know".
Context:
{context}
User's question:
{question}
"""
def load_txt_from_url(url: str) -> Document:
response = requests.get(url)
if response.status_code == 200:
text = response.text.strip()
if not text:
raise ValueError(f"TXT file at {url} is empty.")
return Document(page_content=text, metadata={"source": url})
else:
raise Exception(f"Failed to load {url} with status {response.status_code}")
class RemoteLLM(LLM):
temperature: float = 0.5
top_p: float = 0.95
def __init__(self, temperature: float = 0.5, top_p: float = 0.95):
super().__init__()
self.temperature = temperature
self.top_p = top_p
@property
def _llm_type(self) -> str:
return "remote_llm"
def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
try:
response = requests.post(
"http://localhost:11434/api/generate",
json={
"model": "llama2",
"prompt": prompt,
"temperature": self.temperature,
"top_p": self.top_p
},
stream=False
)
if response.status_code == 200:
return response.json()["response"]
else:
return f"Error: {response.status_code}"
except Exception as e:
return f"Error: {str(e)}"
@property
def _identifying_params(self) -> dict:
return {
"temperature": self.temperature,
"top_p": self.top_p
}
class MistralLLM(LLM):
temperature: float = 0.7
top_p: float = 0.95
_client: Any = PrivateAttr(default=None)
def __init__(self, api_key: str, temperature: float = 0.7, top_p: float = 0.95, **kwargs: Any):
try:
super().__init__(**kwargs)
object.__setattr__(self, '_client', Mistral(api_key=api_key))
self.temperature = temperature
self.top_p = top_p
except Exception as e:
debug_print(f"Init Mistral failed with error: {e}")
@property
def _llm_type(self) -> str:
return "mistral_llm"
def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
try:
debug_print("Calling Mistral API...")
response = self._client.chat.complete(
model="mistral-small-latest",
messages=[{"role": "user", "content": prompt}],
temperature=self.temperature,
top_p=self.top_p
)
return response.choices[0].message.content
except Exception as e:
debug_print(f"Mistral API error: {str(e)}")
return f"Error generating response: {str(e)}"
@property
def _identifying_params(self) -> dict:
return {"model": "mistral-small-latest"}
class LocalLLM(LLM):
@property
def _llm_type(self) -> str:
return "local_llm"
def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
truncated_prompt = truncate_prompt(prompt)
return f"Local LLM Fallback Response for: {truncated_prompt[:100]}..."
@property
def _identifying_params(self) -> dict:
return {}
class ErrorLLM(LLM):
@property
def _llm_type(self) -> str:
return "error_llm"
def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
return "Error: LLM pipeline could not be created. Please check your configuration and try again."
@property
def _identifying_params(self) -> dict:
return {}
class SimpleLLMChain:
def __init__(self, llm_choice: str = "Meta-Llama-3",
temperature: float = 0.5,
top_p: float = 0.95) -> None:
self.llm_choice = llm_choice
self.temperature = temperature
self.top_p = top_p
self.llm = self.create_llm_pipeline()
self.conversation_history = [] # Keep track of conversation
def create_llm_pipeline(self):
from langchain.llms.base import LLM # Import LLM here so it's always defined
normalized = self.llm_choice.lower()
try:
if "remote" in normalized:
debug_print("Creating remote Meta-Llama-3 pipeline via Hugging Face Inference API...")
from huggingface_hub import InferenceClient
repo_id = "meta-llama/Meta-Llama-3-8B-Instruct"
hf_api_token = os.environ.get("HF_API_TOKEN")
if not hf_api_token:
raise ValueError("Please set the HF_API_TOKEN environment variable to use remote inference.")
client = InferenceClient(token=hf_api_token, timeout=120)
# We no longer use wait_for_model because it's unsupported
def remote_generate(prompt: str) -> str:
max_retries = 3
backoff = 2 # start with 2 seconds
for attempt in range(max_retries):
try:
debug_print(f"Remote generation attempt {attempt+1}")
response = client.text_generation(
prompt,
model=repo_id,
temperature=self.temperature,
top_p=self.top_p,
max_new_tokens=512 # Reduced token count for speed
)
return response
except Exception as e:
debug_print(f"Attempt {attempt+1} failed with error: {e}")
if attempt == max_retries - 1:
raise
time.sleep(backoff)
backoff *= 2 # exponential backoff
return "Failed to generate response after multiple attempts."
class RemoteLLM(LLM):
@property
def _llm_type(self) -> str:
return "remote_llm"
def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
return remote_generate(prompt)
@property
def _identifying_params(self) -> dict:
return {"model": repo_id}
debug_print("Remote Meta-Llama-3 pipeline created successfully.")
return RemoteLLM()
elif "mistral" in normalized:
api_key = os.getenv("MISTRAL_API_KEY")
return MistralLLM(api_key=api_key, temperature=self.temperature, top_p=self.top_p)
else:
return LocalLLM()
except Exception as e:
debug_print(f"Error creating LLM pipeline: {str(e)}")
return ErrorLLM()
def update_llm_pipeline(self, new_model_choice: str, temperature: float, top_p: float):
self.llm_choice = new_model_choice
self.temperature = temperature
self.top_p = top_p
self.llm = self.create_llm_pipeline()
def submit_query(self, query: str) -> tuple:
try:
response = self.llm(query)
# Store in conversation history
self.conversation_history.append({"query": query, "response": response})
input_tokens = count_tokens(query)
output_tokens = count_tokens(response)
return (response, f"Input tokens: {input_tokens}", f"Output tokens: {output_tokens}")
except Exception as e:
return (f"Error processing query: {str(e)}", "Input tokens: 0", "Output tokens: 0")
# Update submit_query_updated to work with the simplified chain
def submit_query_updated(query: str, model_choice: str = None, temperature: float = 0.5, top_p: float = 0.95):
"""Process a query with the specified model and parameters."""
debug_print(f"Processing query: {query}")
if not query:
debug_print("Empty query received")
return "Please enter a non-empty query", "", "Input tokens: 0", "Output tokens: 0"
try:
global llm_chain
if llm_chain is None:
llm_chain = SimpleLLMChain(
llm_choice=model_choice,
temperature=temperature,
top_p=top_p
)
elif llm_chain.llm_choice != model_choice:
llm_chain.update_llm_pipeline(model_choice, temperature, top_p)
response, input_tokens, output_tokens = llm_chain.submit_query(query)
return response, "", input_tokens, output_tokens
except Exception as e:
debug_print(f"Error in submit_query_updated: {str(e)}")
return f"Error: {str(e)}", "", "Input tokens: 0", "Output tokens: 0"
# Update the reset_app_updated function
def reset_app_updated():
global llm_chain
llm_chain = None
return "Application reset successfully"
# ----------------------------
# Gradio Interface Functions
# ----------------------------
global rag_chain
rag_chain = SimpleLLMChain()
def load_pdfs_updated(file_links, model_choice, prompt_template, bm25_weight, temperature, top_p):
debug_print("Inside load_pdfs function.")
if not file_links:
debug_print("Please enter non-empty URLs")
return "Please enter non-empty URLs", "Word count: N/A", "Model used: N/A", "Context: N/A"
try:
links = [link.strip() for link in file_links.split("\n") if link.strip()]
global rag_chain
if rag_chain.raw_data:
rag_chain.update_llm_pipeline(model_choice, temperature, top_p, prompt_template, bm25_weight)
context_display = rag_chain.get_current_context()
response_msg = f"Files already loaded. Chain updated with model: {model_choice}"
return (
response_msg,
f"Word count: {word_count(rag_chain.context)}",
f"Model used: {rag_chain.llm_choice}",
f"Context:\n{context_display}"
)
else:
rag_chain = SimpleLLMChain(
llm_choice=model_choice,
temperature=temperature,
top_p=top_p
)
rag_chain.add_pdfs_to_vectore_store(links)
context_display = rag_chain.get_current_context()
response_msg = f"Files loaded successfully. Using model: {model_choice}"
return (
response_msg,
f"Word count: {word_count(rag_chain.context)}",
f"Model used: {rag_chain.llm_choice}",
f"Context:\n{context_display}"
)
except Exception as e:
error_msg = traceback.format_exc()
debug_print("Could not load files. Error: " + error_msg)
return (
"Error loading files: " + str(e),
f"Word count: {word_count('')}",
f"Model used: {rag_chain.llm_choice}",
"Context: N/A"
)
def update_model(new_model: str):
global rag_chain
if rag_chain and rag_chain.raw_data:
rag_chain.update_llm_pipeline(new_model, rag_chain.temperature, rag_chain.top_p,
rag_chain.prompt_template, rag_chain.bm25_weight)
debug_print(f"Model updated to {rag_chain.llm_choice}")
return f"Model updated to: {rag_chain.llm_choice}"
else:
return "No files loaded; please load files first."
def reset_app_updated():
global rag_chain
rag_chain = SimpleLLMChain()
debug_print("App reset successfully.")
return (
"App reset successfully. You can now load new files",
"",
"Model used: Not selected"
)
# ----------------------------
# Gradio Interface Setup
# ----------------------------
custom_css = """
textarea {
overflow-y: scroll !important;
max-height: 200px;
}
"""
# Function to add dots and reset
def add_dots_and_reset():
if not hasattr(add_dots_and_reset, "dots"):
add_dots_and_reset.dots = "" # Initialize the attribute
# Add a dot
add_dots_and_reset.dots += "."
# Reset after 5 dots
if len(add_dots_and_reset.dots) > 5:
add_dots_and_reset.dots = ""
print(f"Current dots: {add_dots_and_reset.dots}") # Debugging print
return add_dots_and_reset.dots
# Define a dummy function to simulate data retrieval
def run_query(max_value):
# Simulate a data retrieval or processing function
return [[i, i**2] for i in range(1, max_value + 1)]
# Function to call both refresh_job_list and check_job_status using the last job ID
def periodic_update(is_checked):
interval = 2 if is_checked else None
debug_print(f"Auto-refresh checkbox is {'checked' if is_checked else 'unchecked'}, every={interval}")
if is_checked:
global last_job_id
job_list_md = refresh_job_list()
job_status = check_job_status(last_job_id) if last_job_id else ("No job ID available", "", "", "", "")
query_results = run_query(10) # Use a fixed value or another logic if needed
return job_list_md, job_status[0], query_results, "" # Return empty string instead of context
else:
# Return empty values to stop updates
return "", "", [], ""
# Define a function to determine the interval based on the checkbox state
def get_interval(is_checked):
return 2 if is_checked else None
# Update the Gradio interface to include job status checking
with gr.Blocks(css=custom_css, js="""
document.addEventListener('DOMContentLoaded', function() {
// Add event listener for job list clicks
const jobListInterval = setInterval(() => {
const jobLinks = document.querySelectorAll('.job-list-container a');
if (jobLinks.length > 0) {
jobLinks.forEach(link => {
link.addEventListener('click', function(e) {
e.preventDefault();
const jobId = this.textContent.split(' ')[0];
// Find the job ID input textbox and set its value
const jobIdInput = document.querySelector('.job-id-input input');
if (jobIdInput) {
jobIdInput.value = jobId;
// Trigger the input event to update Gradio's state
jobIdInput.dispatchEvent(new Event('input', { bubbles: true }));
}
});
});
clearInterval(jobListInterval);
}
}, 500);
});
""") as app:
gr.Markdown('''# PsyLLM Interface
**Model Selection & Parameters:** Choose from the following options:
- ๐บ๐ธ Remote Meta-Llama-3 - has context windows of 8000 tokens
- ๐ช๐บ Mistral-API - has context windows of 32000 tokens
**๐ฅ Randomness (Temperature):** Adjusts output predictability.
- Example: 0.2 makes the output very deterministic (less creative), while 0.8 introduces more variety and spontaneity.
**๐ฏ Word Variety (Topโp):** Limits word choices to a set probability percentage.
- Example: 0.5 restricts output to the most likely 50% of token choices for a focused answer; 0.95 allows almost all possibilities for more diverse responses.
**โ ๏ธ IMPORTANT: This app uses asynchronous processing to avoid timeout issues**
- When you submit a query, you'll receive a Job ID
- Use the "Check Job Status" tab to monitor and retrieve your results
''')
with gr.Tabs() as tabs:
with gr.TabItem("Submit Query"):
with gr.Row():
with gr.Column(scale=1):
llama_checkbox = gr.Checkbox(
value=True,
label="๐บ๐ธ Remote Meta-Llama-3",
info="Context window: 8000 tokens"
)
mistral_checkbox = gr.Checkbox(
value=False,
label="๐ช๐บ Mistral-API",
info="Context window: 32000 tokens"
)
with gr.Column(scale=2):
temperature_slider = gr.Slider(
minimum=0.1, maximum=1.0, value=0.5, step=0.1,
label="Randomness (Temperature)"
)
top_p_slider = gr.Slider(
minimum=0.1, maximum=0.99, value=0.95, step=0.05,
label="Word Variety (Top-p)"
)
with gr.Row():
query_input = gr.Textbox(
label="Enter your query here",
placeholder="Type your query",
lines=4
)
submit_button = gr.Button("Submit Query to Selected Models")
with gr.Row():
with gr.Column(scale=1):
gr.Markdown("### Llama Results")
llama_response = gr.Textbox(
label="Llama Response",
placeholder="Response will appear here",
lines=6
)
llama_tokens = gr.Markdown("Input/Output tokens: 0/0")
with gr.Column(scale=1):
gr.Markdown("### Mistral Results")
mistral_response = gr.Textbox(
label="Mistral Response",
placeholder="Response will appear here",
lines=6
)
mistral_tokens = gr.Markdown("Input/Output tokens: 0/0")
with gr.TabItem("Check Job Status"):
with gr.Row():
with gr.Column(scale=1):
job_list = gr.Markdown(
value="No jobs yet",
label="Job List (Click to select)"
)
# Add the Refresh Job List button
refresh_button = gr.Button("Refresh Job List")
# Use a Checkbox to control the periodic updates
auto_refresh_checkbox = gr.Checkbox(
label="Enable Auto Refresh",
value=False # Default to unchecked
)
# Use a DataFrame to display results
df = gr.DataFrame(
value=run_query(10), # Initial value
headers=["Number", "Square"],
label="Query Results",
visible=False # Set the DataFrame to be invisible
)
with gr.Column(scale=2):
job_id_input = gr.Textbox(
label="Job ID",
placeholder="Job ID will appear here when selected from the list",
lines=1
)
job_query_display = gr.Textbox(
label="Job Query",
placeholder="The query associated with this job will appear here",
lines=2,
interactive=False
)
check_button = gr.Button("Check Status")
cleanup_button = gr.Button("Cleanup Old Jobs")
with gr.Row():
status_response = gr.Textbox(
label="Job Result",
placeholder="Job result will appear here",
lines=6
)
status_context = gr.Textbox(
label="Context Information",
placeholder="Context information will appear here",
lines=6
)
with gr.Row():
status_tokens1 = gr.Markdown("")
status_tokens2 = gr.Markdown("")
with gr.TabItem("App Management"):
with gr.Row():
reset_button = gr.Button("Reset App")
with gr.Row():
reset_response = gr.Textbox(
label="Reset Response",
placeholder="Reset confirmation will appear here",
lines=2
)
reset_context = gr.Textbox(
label="",
placeholder="",
lines=2,
visible=False
)
with gr.Row():
reset_model = gr.Markdown("")
# Connect the buttons to their respective functions
submit_button.click(
submit_query_async,
inputs=[
query_input,
llama_checkbox,
mistral_checkbox,
temperature_slider,
top_p_slider
],
outputs=[
llama_response,
llama_tokens,
mistral_response,
mistral_tokens,
job_id_input,
job_query_display,
job_list
]
)
check_button.click(
check_job_status,
inputs=[job_id_input],
outputs=[status_response, status_context, status_tokens1, status_tokens2, job_query_display]
)
refresh_button.click(
refresh_job_list,
inputs=[],
outputs=[job_list]
)
job_id_input.change(
job_selected,
inputs=[job_id_input],
outputs=[job_id_input, job_query_display]
)
cleanup_button.click(
cleanup_old_jobs,
inputs=[],
outputs=[status_response, status_context, status_tokens1]
)
reset_button.click(
reset_app_updated,
inputs=[],
outputs=[reset_response, reset_context, reset_model]
)
app.load(
fn=refresh_job_list,
inputs=None,
outputs=job_list
)
auto_refresh_checkbox.change(
fn=periodic_update,
inputs=[auto_refresh_checkbox],
outputs=[job_list, status_response, df, status_context],
every=2
)
# Add this with your other global variables
global llm_chain
llm_chain = None
if __name__ == "__main__":
debug_print("Launching Gradio interface.")
app.queue().launch(share=False)
|