awacke1 commited on
Commit
92cb880
·
verified ·
1 Parent(s): db05673

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +28 -32
app.py CHANGED
@@ -7,7 +7,7 @@ import os
7
  TITLE = "✍️ AI Story Outliner"
8
  DESCRIPTION = """
9
  Enter a prompt and get 10 unique story outlines from a CPU-friendly AI model.
10
- The app uses **Tencent's Hunyuan-1.8B** to generate creative outlines formatted in Markdown.
11
 
12
  **How it works:**
13
  1. Enter your story idea.
@@ -25,36 +25,21 @@ examples = [
25
  ]
26
 
27
  # --- Model Initialization ---
28
- # This section loads a smaller, CPU-friendly model.
29
- # It will automatically use the HF_TOKEN secret when deployed on Hugging Face Spaces.
30
  generator = None
31
  model_error = None
32
 
33
  try:
34
  print("Initializing model... This may take a moment.")
35
 
36
- # Explicitly load the token from environment variables (for HF Spaces secrets).
37
- # This makes the authentication more robust, overriding any bad default credentials.
38
- hf_token = os.environ.get("HF_TOKEN")
39
-
40
- # Add a check to see if the token was loaded correctly.
41
- if hf_token:
42
- print("✅ HF_TOKEN secret found.")
43
- else:
44
- print("⚠️ HF_TOKEN secret not found. Please ensure it is set in your Hugging Face Space settings.")
45
- # Raise an error to stop the app from proceeding without a token.
46
- raise ValueError("Hugging Face token not found. Please set the HF_TOKEN secret.")
47
-
48
- # Using a smaller model from the user's list.
49
- # Passing the token explicitly to ensure correct authentication.
50
  generator = pipeline(
51
  "text-generation",
52
- model="tencent/Hunyuan-1.8B-Instruct",
53
- torch_dtype=torch.bfloat16, # Use bfloat16 for better performance if available
54
- device_map="auto", # Will use GPU if available, otherwise CPU
55
- token=hf_token
56
  )
57
- print("✅ Tencent/Hunyuan-1.8B-Instruct model loaded successfully!")
58
 
59
  except Exception as e:
60
  model_error = e
@@ -76,31 +61,42 @@ def generate_stories(prompt: str) -> list[str]:
76
  # Return a list of 10 empty strings to clear the outputs
77
  return [""] * 10
78
 
79
- # This prompt format is specific to the Hunyuan model.
80
- system_instruction = "You are an expert storyteller. Your task is to take a user's prompt and write a short story as a Markdown outline. The story must have a dramatic arc and be the length of a song. Use emojis to highlight the story sections."
81
-
82
- story_prompt = f"<|im_start|>system\n{system_instruction}<|im_end|>\n<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
 
 
 
 
 
 
 
 
 
 
 
 
83
 
84
  # Parameters for the pipeline to generate 10 diverse results.
85
  params = {
86
  "max_new_tokens": 250,
87
  "num_return_sequences": 10,
88
  "do_sample": True,
89
- "temperature": 0.8,
 
90
  "top_p": 0.95,
91
- "pad_token_id": generator.tokenizer.eos_token_id # Suppress warning
92
  }
93
 
94
  # Generate 10 different story variations
95
  outputs = generator(story_prompt, **params)
96
 
97
- # Extract the generated text and clean it up.
98
  stories = []
99
  for out in outputs:
100
- # Remove the prompt part from the full generated text
101
  full_text = out['generated_text']
102
- assistant_response = full_text.split("<|im_start|>assistant\n")[-1]
103
- stories.append(assistant_response)
104
 
105
  # Ensure we return exactly 10 stories, padding if necessary.
106
  while len(stories) < 10:
 
7
  TITLE = "✍️ AI Story Outliner"
8
  DESCRIPTION = """
9
  Enter a prompt and get 10 unique story outlines from a CPU-friendly AI model.
10
+ The app uses **GPT-2**, a reliable and well-established model, to generate creative outlines.
11
 
12
  **How it works:**
13
  1. Enter your story idea.
 
25
  ]
26
 
27
  # --- Model Initialization ---
28
+ # This section loads a smaller, stable, and CPU-friendly model that requires no authentication.
 
29
  generator = None
30
  model_error = None
31
 
32
  try:
33
  print("Initializing model... This may take a moment.")
34
 
35
+ # Using 'gpt2', a stable and widely supported model that does not require a token.
 
 
 
 
 
 
 
 
 
 
 
 
 
36
  generator = pipeline(
37
  "text-generation",
38
+ model="gpt2",
39
+ torch_dtype=torch.float32, # Use float32 for wider CPU compatibility
40
+ device_map="auto" # Will use GPU if available, otherwise CPU
 
41
  )
42
+ print("✅ GPT-2 model loaded successfully!")
43
 
44
  except Exception as e:
45
  model_error = e
 
61
  # Return a list of 10 empty strings to clear the outputs
62
  return [""] * 10
63
 
64
+ # A generic story prompt that works well with models like GPT-2.
65
+ story_prompt = f"""
66
+ Story Idea: "{prompt}"
67
+
68
+ Create a short story outline based on this idea.
69
+
70
+ ### 🎬 The Hook
71
+ A dramatic opening.
72
+
73
+ ### 🎼 The Ballad
74
+ The main story, told concisely.
75
+
76
+ ### 🔚 The Finale
77
+ A clear and satisfying ending.
78
+ ---
79
+ """
80
 
81
  # Parameters for the pipeline to generate 10 diverse results.
82
  params = {
83
  "max_new_tokens": 250,
84
  "num_return_sequences": 10,
85
  "do_sample": True,
86
+ "temperature": 0.9,
87
+ "top_k": 50,
88
  "top_p": 0.95,
89
+ "pad_token_id": generator.tokenizer.eos_token_id
90
  }
91
 
92
  # Generate 10 different story variations
93
  outputs = generator(story_prompt, **params)
94
 
95
+ # Extract the generated text. GPT-2 will continue from the prompt.
96
  stories = []
97
  for out in outputs:
 
98
  full_text = out['generated_text']
99
+ stories.append(full_text)
 
100
 
101
  # Ensure we return exactly 10 stories, padding if necessary.
102
  while len(stories) < 10: