Ilneval/DarkFantasyWritingAI

1import gradio as gr2from transformers import AutoModelForCausalLM, AutoTokenizer3 4# Load GPT-J model5model_name = "EleutherAI/gpt-j-6B"6tokenizer = AutoTokenizer.from_pretrained(model_name)7model = AutoModelForCausalLM.from_pretrained(model_name)8 9# Function for generating responses10def generate_response(prompt):11    inputs = tokenizer(prompt, return_tensors="pt", padding=True)12    outputs = model.generate(**inputs, max_length=500)13    response = tokenizer.decode(outputs[0], skip_special_tokens=True)14    return response15 16# Custom AI Assistant UI17custom_instructions = """18You are a writing assistant specialized in helping create a dark fantasy novel based on a Dungeons & Dragons campaign. 19- Maintain a classic high fantasy style with immersive descriptions and deep worldbuilding.20- Ensure character dialogue is engaging and stays true to the provided backstories.21- Help structure novel chapters based on session recaps.22- The world should feel dangerous yet filled with moments of camaraderie.23- Avoid modern language or excessive exposition.24"""25 26def chat(input_text):27    full_prompt = f"{custom_instructions}\n\nUser: {input_text}\nAssistant:"28    return generate_response(full_prompt)29 30# Gradio UI31iface = gr.Interface(fn=chat, inputs="text", outputs="text", title="Dark Fantasy Writing AI")32iface.launch()