persona_chatbot.py

'''from datasets import load_dataset
from llama_cpp import Llama

model_path = "path/to/your/model.gguf"


# Or from HF Hub:
dataset = load_dataset("K00B404/personas")

personas = dataset["persona"]
unique_personas = list(set(personas))

print("Available personas:")
for i, p in enumerate(unique_personas):
    print(f"{i}: {p}")

choice = int(input("Select a persona by number: "))
selected_persona = unique_personas[choice]
print(f"Selected persona: {selected_persona}")


# Find the row(s) for the selected persona
persona_info = dataset.filter(lambda x: x["persona"] == selected_persona)[0]

system_prompt = f"""You are a chatbot with the following persona:
Persona: {persona_info['persona']}
Sex: {persona_info['sex']}
Age: {persona_info['age']}

Please respond accordingly.
"""

llm = Llama(model_path=model_path)

def chat_with_llm(system_prompt):
    print("Chatbot is ready! Type 'exit' to quit.")
    while True:
        user_input = input("You: ")
        if user_input.lower() == "exit":
            break

        prompt = f"{system_prompt}\nUser: {user_input}\nAssistant:"
        response = llm(prompt=prompt, max_tokens=256)
        print("Bot:", response['choices'][0]['text'].strip())

chat_with_llm(system_prompt)
'''

from datasets import load_dataset
from llama_cpp import Llama

class PersonaChatbot:
    def __init__(self, model_path: str, dataset_name: str):
        self.model_path = model_path
        self.dataset_name = dataset_name
        
        # Load dataset
        print(f"Loading dataset '{dataset_name}'...")
        self.dataset = load_dataset(dataset_name, split="train")
        
        # Extract unique personas
        self.unique_personas = list(set(self.dataset["persona"]))
        
        # Load model
        print(f"Loading LLM model from '{model_path}'...")
        self.llm = Llama(model_path=model_path)
        
        self.system_prompt = ""
        
    def list_personas(self):
        print("Available personas:")
        for i, p in enumerate(self.unique_personas):
            print(f"{i}: {p}")
    
    def select_persona(self, index: int):
        if index < 0 or index >= len(self.unique_personas):
            raise ValueError("Persona index out of range")
        selected = self.unique_personas[index]
        print(f"Selected persona: {selected}")
        
        # Filter dataset rows for selected persona and pick the first one
        filtered = self.dataset.filter(lambda x: x["persona"] == selected)
        if len(filtered) == 0:
            raise ValueError("No data found for selected persona")
        
        persona_info = filtered[0]
        self.system_prompt = (
            f"You are a chatbot with the following persona:\n"
            f"Persona: {persona_info['persona']}\n"
            f"Sex: {persona_info['sex']}\n"
            f"Age: {persona_info['age']}\n"
            f"Please respond accordingly."
        )
    
    def chat(self):
        if not self.system_prompt:
            print("You must select a persona first using select_persona()")
            return
        
        print("Chatbot is ready! Type 'exit' to quit.")
        while True:
            user_input = input("You: ")
            if user_input.lower() == "exit":
                print("Goodbye!")
                break
            
            prompt = f"{self.system_prompt}\nUser: {user_input}\nAssistant:"
            response = self.llm(prompt=prompt, max_tokens=256)
            print("Bot:", response['choices'][0]['text'].strip())

# Usage example:

if __name__ == "__main__":
    MODEL_PATH = "path/to/your/model.gguf"
    DATASET_NAME = "K00B404/personas"
    
    bot = PersonaChatbot(MODEL_PATH, DATASET_NAME)
    bot.list_personas()
    choice = int(input("Select a persona by number: "))
    bot.select_persona(choice)
    bot.chat()