Update app.py

app.py

@@ -1,6 +1,29 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import gradio as gr
 
-def load_model():
-    model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
-    tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
-    return model, tokenizer
+def initialize_model():
+    model = AutoModelForCausalLM.from_pretrained("BeaverAI/mistral-doryV2-12b")
+    tokenizer = AutoTokenizer.from_pretrained("BeaverAI/mistral-doryV2-12b")
+    return model, tokenizer
+
+def generate_story(prompt):
+    model, tokenizer = initialize_model()
+    inputs = tokenizer(f"""### System:
+    You are a creative story assistant.
+    
+### Instruction:
+{prompt}
+
+### Response:
+""", return_tensors="pt")
+    
+    outputs = model.generate(inputs["input_ids"], max_length=1000)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+interface = gr.Interface(
+    fn=generate_story,
+    inputs="text",
+    outputs="text"
+)
+
+interface.launch()