Update app.py

app.py

@@ -1,8 +1,8 @@
 import os
 from threading import Thread
 from typing import Iterator
-import json
-import uuid
+
+from mongoengine import connect, Document, StringField, SequenceField
 
 import gradio as gr
 import spaces
@@ -41,9 +41,17 @@ if torch.cuda.is_available():
     model = PeftModel.from_pretrained(base_model,"ranamhamoud/storytell")
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     tokenizer.pad_token = tokenizer.eos_token
+
+PASSWORD = os.environ.get("MONGODB_PASS")
+connect(host='mongodb+srv://ranamhammoud11:{PASSWORD}@stories.zf5v52a.mongodb.net/')
+
+class Story(Document):
+    message = StringField()
+    content = StringField()
+    story_id = SequenceField(primary_key=True)
     
 def make_prompt(entry):
-    return f"### Human: YOUR INSTRUCTION HERE,ONLY TELL A STORY: {entry} ### Assistant:"
+    return f"### Human: YOUR INSTRUCTION HERE,ALWAYS USE A STORY,INCLUDE ASSESMENTS,TECHNICAL SUMMARY: {entry} ### Assistant:"
     
 @spaces.GPU
 def generate(
@@ -88,6 +96,11 @@ def generate(
     for text in streamer:
         outputs.append(text)
         yield "".join(outputs)
+    try:
+        saved_story = Story(message=message, content=final_story).save() 
+        yield f"Story saved with ID: {saved_story.story_id}"
+    except Exception as e:
+        yield f"Failed to save story: {str(e)}"
 
 chat_interface = gr.ChatInterface(
     fn=generate,