Update app.py

app.py

@@ -61,8 +61,8 @@ def text2story(text):
         genre = labels[predicted_class]
         
         # Then, generate a story based on the predicted genre
-        story_generator = AutoModelForCausalLM.from_pretrained("gpt2")
-        story_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        story_generator = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
+        story_tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
         
         # Create a genre-specific prompt
         prompt = f"Write a {genre.lower()} story about: {text}\n\nStory:"
@@ -85,6 +85,25 @@ def text2story(text):
         story = story_tokenizer.decode(outputs[0], skip_special_tokens=True)
         final_story = story.replace(prompt, "").strip()
         return final_story
+    except Exception as e:
+        # just Fallback only used  openai-community/gpt2 if the advanced one fails
+        fallback_generator = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
+        fallback_prompt = f"{text}"
+        tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
+        inputs = tokenizer(fallback_prompt, return_tensors="pt")
+        fallback_story = fallback_generator.generate(
+            inputs.input_ids,
+            min_length=50,
+            max_new_tokens=100,
+            temperature=0.7,
+            top_p=0.9,
+            top_k=40,
+            repetition_penalty=1.2,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        fallback_story = tokenizer.decode(fallback_story[0], skip_special_tokens=True)
+        return fallback_story
 
 def load_css(css_file):
     with open(css_file) as f: