Update app.py

app.py

@@ -12,25 +12,44 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype="auto"
 )
 
+# Ensure pad token is set for safe generation
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+if getattr(model.config, "pad_token_id", None) is None:
+    model.config.pad_token_id = tokenizer.pad_token_id
+    
 # Create generation pipeline
-story_generator = pipeline(
-    "text-generation",
+#story_generator = pipeline(
+    #"text-generation",
+    #model=model,
+    #tokenizer=tokenizer )
+
+# Build a text generation pipeline
+generator = pipeline(
+    task="text-generation",
     model=model,
-    tokenizer=tokenizer
+    tokenizer=tokenizer,
+    return_full_text=False
 )
 
 # Function to generate stories
-def generate_story(prompt, max_length=300, temperature=0.8):
-    outputs = story_generator(
-        model,
-        #prompt,
-        max_length=max_length,
-        temperature=temperature,
-        do_sample=True,
-        top_p=0.95,
-        top_k=50
-    )
-    return outputs[0]["generated_text"]
+def generate_story(prompt, max_tokens=300, temperature=0.8):
+    try:
+        chat_prompt = to_chat_prompt(prompt)
+        outputs = generator(
+            chat_prompt,
+            max_new_tokens=int(max_tokens),
+            temperature=float(temperature),
+            do_sample=True,
+            top_p=0.95,
+            top_k=50,
+            repetition_penalty=1.05,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id
+        )
+        return outputs[0]["generated_text"]
+    except Exception as e:
+        return f"Error during generation: {type(e).__name__}: {e}"
 
 # Gradio UI
 with gr.Blocks() as demo:
@@ -42,7 +61,7 @@ with gr.Blocks() as demo:
         lines=3
     )
 
-    max_length = gr.Slider(50, 1000, value=300, step=50, label="Story Length")
+    max_length = gr.Slider(50, 1000, value=300, step=50, label="Story Length in new tokens")
     temperature = gr.Slider(0.1, 1.5, value=0.8, step=0.1, label="Creativity")
 
     generate_btn = gr.Button("✨ Generate Story")