app.py

import transformers
import torch
import gradio as gr

model_id = "yodayo-ai/nephra_v1.0"

pipeline = transformers.pipeline(
    "text-generation",
    model=model_id,
    model_kwargs={"torch_dtype": torch.bfloat16},
    device_map="auto",
    offload_folder="offload",  # Add this line to specify the offload folder
)

# Define characters
characters = [
    {
        "name": "Alex",
        "description": "Alex is a young and ambitious adventurer, full of energy and eager to discover new places. He is always ready to face any challenges and strives to explore the unknown.",
        "traits": "brave, energetic, optimistic, determined"
    },
    {
        "name": "Maya",
        "description": "Maya is a wise and experienced sorceress, possessing deep knowledge of magic and ancient rituals. She is known for her calm demeanor, analytical mind, and ability to find solutions in difficult situations.",
        "traits": "calm, thoughtful, intuitive, attentive"
    },
    {
        "name": "Victor",
        "description": "Victor is a former warrior who left behind his fighting days to seek inner peace and harmony. His life experience and sense of justice make him a reliable friend and mentor.",
        "traits": "serious, judicious, fair, balanced"
    }
]

# Function to generate response
def generate_response(character_name, user_input, max_length=100, temperature=1.12, min_p=0.075, repetition_penalty=1.1):
    # Find the character
    character = next((c for c in characters if c["name"] == character_name), None)
    if not character:
        return "Character not found."

    # Prepare the message based on the selected character's personality and description
    messages = [
        {"role": "system", "content": f"You are {character_name}, {character['description']}. Personality traits: {character['traits']}."},
        {"role": "user", "content": user_input},
    ]

    # Prepare the prompt using the chat template from the pipeline's tokenizer
    prompt = pipeline.tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    # Generate response using the pipeline
    outputs = pipeline(
        prompt,
        max_new_tokens=max_length,
        eos_token_id=[
            pipeline.tokenizer.convert_tokens_to_ids(""),
            pipeline.tokenizer.eos_token_id,
        ],
        do_sample=True,
        temperature=temperature,
        min_p=min_p,
        repetition_penalty=repetition_penalty,
    )

    # Extract and return the generated response
    generated_text = outputs[0]["generated_text"][len(prompt):].strip()
    return generated_text

# Gradio Interface
iface = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Dropdown([c["name"] for c in characters], label="Choose a character"),
        gr.Textbox(lines=2, placeholder="Enter your text here..."),
        gr.Slider(20, 200, step=1, value=100, label="Max Length"),
        gr.Slider(0.1, 1.0, step=0.1, value=1.12, label="Temperature"),
        gr.Slider(0.01, 1.0, step=0.01, value=0.075, label="min-p"),
        gr.Slider(1.0, 2.0, step=0.1, value=1.1, label="Repetition Penalty")
    ],
    outputs="text",
    title="Nephra v1 LLM Roleplaying Chatbot",
    description="Enter a text prompt to generate a response using the Nephra v1 model, based on the selected character."
)

if __name__ == "__main__":
    iface.launch(share=True, debug=True)