Spaces:
Sleeping
Sleeping
| 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) |