Skip to content

Update cli.py #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shouryashashank wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Update cli.py #5

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0509f4f
Update cli.py
shouryashashank Oct 22, 2024
7c037b9
Update cli.py
shouryashashank Nov 2, 2024
f582a12
Update cli.py
shouryashashank Nov 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 39 additions & 8 deletions app/predacons_cli/src/cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from rich.table import Table
from rich.markdown import Markdown
from rich.console import Console
from transformers import TextStreamer
try:
from .rag import VectorStore
from .rag import WebScraper
Expand Down Expand Up @@ -97,6 +98,10 @@ def launch(self,logs=False):
for i in range(1, 0, -1):
time.sleep(1)
os.system('clear') # Clear the screen
personality = None
if config.get("personality", None):
print("[yellow]Personality detected! Loading the quirks and traits... Brace yourself![/yellow]")
personality = "\n Make sure to respond back as a " + config["personality"]
Comment on lines +101 to +104

@coderabbitai coderabbitai Bot Oct 22, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Add input validation for personality configuration.

The personality value is used directly in prompt templates without validation. This could lead to prompt injection vulnerabilities or unexpected behavior.

Consider adding validation:

 personality = None
 if config.get("personality", None):
+    # Validate personality against allowed values
+    allowed_personalities = ["friendly", "professional", "casual"]  # Define appropriate values
+    if config["personality"].lower() not in allowed_personalities:
+        print("[red]Warning: Invalid personality specified. Using default.[/red]")
+        personality = None
+    else:
         print("[yellow]Personality detected! Loading the quirks and traits... Brace yourself![/yellow]")
-        personality = "\n Make sure to respond back as a " + config["personality"]
+        personality = f"\nMake sure to respond back in a {config['personality']} manner"
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
personality = None
if config.get("personality", None):
print("[yellow]Personality detected! Loading the quirks and traits... Brace yourself![/yellow]")
personality = "\n Make sure to respond back as a " + config["personality"]
personality = None
if config.get("personality", None):
# Validate personality against allowed values
allowed_personalities = ["friendly", "professional", "casual"] # Define appropriate values
if config["personality"].lower() not in allowed_personalities:
print("[red]Warning: Invalid personality specified. Using default.[/red]")
personality = None
else:
print("[yellow]Personality detected! Loading the quirks and traits... Brace yourself![/yellow]")
personality = f"\nMake sure to respond back in a {config['personality']} manner"


print("[i]Welcome to the Predacons CLI![/i] [green]Model: [orange1]"+config["model_path"]+"[/orange1] loaded successfully![/green]")
print("[yellow]You can start chatting with Predacons now.Type 'clear' to clear history, Type 'exit' to quit, Type 'help' for more options, Type 'update' to update the load documents[/yellow]")
Expand Down Expand Up @@ -143,7 +148,7 @@ def launch(self,logs=False):

Answer the question based on the above context: {question}
"""

PROMPT_TEMPLATE = PROMPT_TEMPLATE + personality if personality else ""

@coderabbitai coderabbitai Bot Oct 22, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Ensure consistent personality integration across all prompt templates.

The personality trait is only appended to the first prompt template when using web search results. It should be consistently applied to all prompt templates for uniform behavior.

Apply this change to all prompt templates:

 PROMPT_TEMPLATE = PROMPT_TEMPLATE + personality if personality else ""

Also applies to: 164-164, 179-179

user_input = PROMPT_TEMPLATE.format(db_context=context_text,web_context=web_text, question=user_input)

PROMPT_TEMPLATE = """
Expand Down Expand Up @@ -173,15 +178,25 @@ def launch(self,logs=False):
user_input = PROMPT_TEMPLATE.format(context=web_text, question=user_input)
user_body = {"role": "user", "content": user_input}
chat.append(user_body)
response = Cli.generate_response(self, chat, model, tokenizer, config)
thread,streamer = Cli.generate_response(self, chat, model, tokenizer, config)
thread.start()
print("[orange1]Predacons: [/orange1]", end="")
try:
response = ""
for new_text in streamer:
response = response + new_text
print("[sky_blue1]"+ new_text + "[/sky_blue1]", end="")
print("\n")
finally:
thread.join()
Comment on lines +181 to +191

@coderabbitai coderabbitai Bot Nov 2, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Ensure thread safety in response generation.

The streaming implementation uses threading but lacks proper thread safety mechanisms. Consider adding thread synchronization to prevent potential race conditions.

Apply this change:

 thread,streamer = Cli.generate_response(self, chat, model, tokenizer, config)
+response_lock = threading.Lock()
 thread.start()
 print("[orange1]Predacons: [/orange1]", end="")
 try:
     response = ""
     for new_text in streamer:
+        with response_lock:
             response = response + new_text
         print("[sky_blue1]"+ new_text + "[/sky_blue1]", end="")
     print("\n")
 finally:
     thread.join()

Don't forget to add the import at the top:

+import threading

Committable suggestion skipped: line range outside the PR's diff.

response_body = {"role": "assistant", "content": response}
chat.append(response_body)
if config["print_as_markdown"]:
markdown = Markdown(response)
print("[orange1]Predacons: [/orange1]")
console.print(markdown)
else:
console.print("[orange1]Predacons: [/orange1] [sky_blue1]" + response+"[/sky_blue1]")
# else:
# console.print("[orange1]Predacons: [/orange1] [sky_blue1]" + response+"[/sky_blue1]")


def load_model(self, model_path,trust_remote_code=False,use_fast_generation=False, draft_model_name=None,gguf_file=None,auto_quantize=None):
Expand Down Expand Up @@ -269,7 +284,8 @@ def create_config_file(self):
"vector_db_path": None,
"document_path": None,
"embedding_model" : None,
"scrap_web": False
"scrap_web": False,
"personality": None
}

# If no config file is found, use the default configuration
Expand Down Expand Up @@ -330,6 +346,7 @@ def create_config_file(self):
document_path = Prompt.ask("Enter the document path", default=config["document_path"])
embedding_model = Prompt.ask("Enter the embedding model", default=config["embedding_model"])
scrap_web = Prompt.ask("Scrap the web for data? (true/false)", default=str(config["scrap_web"]))
personality = Prompt.ask("Enter the personality", default=config["personality"])


config_data = {
Expand All @@ -350,7 +367,8 @@ def create_config_file(self):
"vector_db_path": vector_db_path if vector_db_path else None,
"document_path": document_path if document_path else None,
"embedding_model": embedding_model if embedding_model else None,
"scrap_web": scrap_web.lower() == 'true'
"scrap_web": scrap_web.lower() == 'true',
"personality": personality if personality else None
}

with open(file_path, 'w') as f:
Expand All @@ -366,6 +384,19 @@ def load_config_file(self):
return config

def generate_response(self, chat, model, tokenizer, config):
thread,streamer = self.predacons.chat_generate(model = model,
sequence = chat,
max_length = config["max_length"],
tokenizer = tokenizer,
trust_remote_code = config["trust_remote_code"],
do_sample=True,
temperature = config["temperature"],
dont_print_output = True,
stream = True
)
return thread,streamer
Comment on lines +387 to +397

@coderabbitai coderabbitai Bot Nov 2, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Remove duplicate response generation method.

The generate_response2 method appears to be a duplicate of generate_response with only the streaming parameter different.

Consider consolidating these methods:

-def generate_response2(self, chat, model, tokenizer, config):
-    response = self.predacons.chat_generate(model = model,
-        sequence = chat,
-        max_length = config["max_length"],
-        tokenizer = tokenizer,
-        trust_remote_code = config["trust_remote_code"],
-        do_sample=True,   
-        temperature = config["temperature"],
-        dont_print_output = True,
-        )
-    return response

 def generate_response(self, chat, model, tokenizer, config, stream=True):
     thread,streamer = self.predacons.chat_generate(model = model,
         sequence = chat,
         max_length = config["max_length"],
         tokenizer = tokenizer,
         trust_remote_code = config["trust_remote_code"],
         do_sample=True,   
         temperature = config["temperature"],
         dont_print_output = True,
         stream = stream
         )
-    return thread,streamer
+    if stream:
+        return thread,streamer
+    return streamer

Also applies to: 398-409


def generate_response2(self, chat, model, tokenizer, config):
response = self.predacons.chat_generate(model = model,
sequence = chat,
max_length = config["max_length"],
Expand All @@ -377,7 +408,7 @@ def generate_response(self, chat, model, tokenizer, config):
)
return response

# cli = Cli()
# cli.launch()
cli = Cli()
cli.launch()
Comment on lines +411 to +412

@coderabbitai coderabbitai Bot Oct 22, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Add proper entry point guard.

Direct script execution should be guarded to prevent unintended execution when the module is imported.

Wrap the execution in a main guard:

-cli = Cli()
-cli.launch()
+if __name__ == "__main__":
+    cli = Cli()
+    cli.launch()
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
cli = Cli()
cli.launch()
if __name__ == "__main__":
cli = Cli()
cli.launch()