Update app.py

app.py

@@ -1,14 +1,21 @@
+import os
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
+# Model configuration
 MODEL_NAME = "OpenAssistant/oasst-sft-1-pythia-12b"
+OFFLOAD_DIR = "./offload"
 
-# Load model & tokenizer
+# Ensure offload directory exists
+os.makedirs(OFFLOAD_DIR, exist_ok=True)
+
+# Load tokenizer and model safely for large weights
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
     device_map="auto",
+    offload_folder=OFFLOAD_DIR,
     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
 )
 model.eval()
@@ -21,9 +28,6 @@ You are a professional ghostwriter. Your writing is:
 - Explicit within acceptable content guidelines
 - Italicized for internal monologue or emotional emphasis
 - Structured and catchy
-- Explicit and authentic
-- Sexual and mixture of sfw and nsfw
-- Realistic and slice of life
 """
 
 def stream_generate(prompt, temperature, max_tokens):
@@ -48,6 +52,7 @@ def stream_generate(prompt, temperature, max_tokens):
         if tokenizer.decode(next_token[0]) in [tokenizer.eos_token, "\n\n"]:
             break
 
+# Gradio interface
 with gr.Blocks(title="🧠 HuggingChat Stream Writer") as demo:
     gr.Markdown("## ✍️ Real-Time HuggingChat-Style Generator")
     gr.Markdown("*Watch your story unfold word by word...*")