Update app.py

app.py

@@ -3,23 +3,28 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
 # Model setup
-#model_name = "EleutherAI/gpt-neo-125M"  # Lightweight model for Spaces
-model_name = "EleutherAI/gpt-neo-1.3B" # A bit better model
+model_name = "EleutherAI/gpt-neo-1.3B" 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name).to("cpu")  # CPU-friendly for Spaces
 
 # Text generation function
 def generate_text(prompt, max_length=500, temperature=0.7, top_p=0.9):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(
-        **inputs,
-        max_length=max_length,
-        temperature=temperature,
-        top_p=top_p,
-        do_sample=True,
-        pad_token_id=tokenizer.eos_token_id
-    )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    try:
+        # Calculate max tokens for the input to fit within model limits
+        max_input_length = 2048 - max_length  # Total context limit minus the max length of output
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=max_input_length)
+        
+        outputs = model.generate(
+            **inputs,
+            max_length=max_length,
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    except Exception as e:
+        return f"Error during generation: {str(e)}"
 
 # Global variables to store hierarchical content
 global_synopsis = ""
@@ -28,30 +33,39 @@ global_chapters = ""
 # Generation functions
 def generate_synopsis(topic):
     global global_synopsis
-    prompt = f"Write a brief synopsis for a story about {topic}: . The max number of characters must be 500"
-    global_synopsis = generate_text(prompt, max_length=500)
-    return global_synopsis
+    try:
+        prompt = f"Write a brief synopsis for a story about {topic}. The max number of characters must be 500."
+        global_synopsis = generate_text(prompt, max_length=500)
+        return global_synopsis
+    except Exception as e:
+        return f"Error generating synopsis: {str(e)}"
 
 def generate_chapters():
     global global_synopsis, global_chapters
     if not global_synopsis:
         return "Please generate a synopsis first."
-    prompt = f'''Based on this synopsis for a book: {global_synopsis}. Divide the story into 4 chapters with brief descriptions for each. 
-    Enumerate every chapter created followed by its description and make the first chapter sound like an introduction and the last may sound as the epilogue.
-    The max number of character for every pair title, description must be 500'''
-    global_chapters = generate_text(prompt, max_length=2000)
-    return global_chapters
+    try:
+        prompt = f'''Based on this synopsis for a book: {global_synopsis}. Divide the story into 4 chapters with brief descriptions for each. 
+        Enumerate every chapter created followed by its description and make the first chapter sound like an introduction and the last may sound as the epilogue.
+        The max number of characters for every pair title, description must be 500.'''
+        global_chapters = generate_text(prompt, max_length=2000)
+        return global_chapters
+    except Exception as e:
+        return f"Error generating chapters: {str(e)}"
 
 def expand_chapter(chapter_number):
     global global_chapters
     if not global_chapters:
         return "Please generate chapters first."
-    chapters = global_chapters.split("\n")
-    if chapter_number <= 0 or chapter_number > len(chapters):
-        return f"Select a number between 1 and {len(chapters)}."
-    prompt = f'''Knowing this synopsis for a book: {global_synopsis}. Expand and describe Chapter {chapter_number} 
-    in more detail, the title and current brief description of this chapter is: {chapters[chapter_number - 1]}'''
-    return generate_text(prompt, max_length=500)
+    try:
+        chapters = global_chapters.split("\n")
+        if chapter_number <= 0 or chapter_number > len(chapters):
+            return f"Select a number between 1 and {len(chapters)}."
+        prompt = f'''Knowing this synopsis for a book: {global_synopsis}. Expand and describe Chapter {chapter_number} 
+        in more detail, the title and current brief description of this chapter is: {chapters[chapter_number - 1]}'''
+        return generate_text(prompt, max_length=500)
+    except Exception as e:
+        return f"Error expanding chapter: {str(e)}"
 
 # Gradio interface
 with gr.Blocks() as demo:
@@ -77,4 +91,4 @@ with gr.Blocks() as demo:
     chapter_button.click(expand_chapter, inputs=chapter_input, outputs=chapter_detail_output)
 
 # Launch the app
-demo.launch()
+demo.launch()