Update app.py

app.py

@@ -44,41 +44,43 @@ def translate_to_chinese(text):
     return translation
 
 
-# text2story - using mosaicml/mpt-7b-storywriter model for better stories
+# text2story - using llama-3-8b-Instruct-bnb-4bit-personal-shortstory model for better stories
 def text2story(text):
     try:
-        # Initialize the improved story generation pipeline
-        generator = pipeline("text-generation", model="2173ars/llama-3-8b-Instruct-bnb-4bit-personal-shortstory")
-        
-        # Create a prompt for the story
+        model = AutoModel.from_pretrained("2173ars/llama-3-8b-Instruct-bnb-4bit-personal-shortstory")
         prompt = f"{text}"
-        
-        story = generator(prompt,
-                         min_length=100,
-                         max_length=130,  
-                         num_return_sequences=1,
-                         top_k=50,  
-                         top_p=0.92, 
-                         no_repeat_ngram_size=3, 
-                         temperature=0.5,
-                         repetition_penalty=1.3)[0]['generated_text']
-        
-        return story.replace(prompt, "").strip()        
-    
+        tokenizer = AutoTokenizer.from_pretrained("2173ars/llama-3-8b-Instruct-bnb-4bit-personal-shortstory")
+        inputs = tokenizer(prompt, return_tensors="pt")
+        outputs = model.generate(
+            inputs.input_ids,
+            max_new_tokens=250,
+            temperature=0.7,
+            top_p=0.9,
+            top_k=40,
+            repetition_penalty=1.2,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        story = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return story  
     except Exception as e:
         # Fallback to simpler model if the advanced one fails
-        fallback_generator = pipeline('text-generation', model='gpt2')
-        fallback_prompt = f"{text} "
-        fallback_story = fallback_generator(fallback_prompt,
-                         min_length=100, 
-                         max_length=130,  
-                         num_return_sequences=1,
-                         top_k=50,  
-                         top_p=0.92, 
-                         no_repeat_ngram_size=3, 
-                         temperature=0.5,
-                         repetition_penalty=1.3)[0]['generated_text']
-        return fallback_story.replace(fallback_prompt, "").strip()
+        fallback_generator = AutoModel.from_pretrained("openai-community/gpt2")
+        fallback_prompt = f"{text}"
+        tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
+        inputs = tokenizer(prompt, return_tensors="pt")
+        fallback_story = model.generate(
+            inputs.input_ids,
+            max_new_tokens=250,
+            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(outputs[0], skip_special_tokens=True)
+        return fallback_story
 
 def load_css(css_file):
     with open(css_file) as f: