Update app.py

app.py

@@ -3,17 +3,19 @@ from transformers import pipeline
 import torch
 import subprocess
 import spaces
+import os
 
 
 # Install flash-attn
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 # Initialize the model pipeline
-generator = pipeline('text-generation', model='mistralai/Mistral-7B-v0.1', torch_dtype=torch.bfloat16)
+generator = pipeline('text-generation', model='Locutusque/Hyperion-1.5-Mistral-7B', torch_dtype=torch.bfloat16, token=os.environ["HF"])
 @spaces.GPU
 def generate_text(prompt, temperature, top_p, top_k, repetition_penalty, max_length):
     # Generate text using the model
     generator.model.cuda()
     generator.device = torch.device("cuda")
+    prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
     outputs = generator(
         prompt, 
         do_sample=True,
@@ -41,8 +43,8 @@ iface = gr.Interface(
         gr.Slider(minimum=5, maximum=4096, step=5, value=1024, label="Max Length")
     ],
     outputs=gr.Textbox(label="Generated Text"),
-    title="Text Completion Model",
-    description="Try out the Mistral-7B model for free! Note this is the pretrained model and is not fine-tuned for instruction."
+    title="Hyperion-1.5-Mistral-7B",
+    description="Try out the Hyperion-1.5-Mistral-7B model for free! This is a preview version, and the model will be released soon"
 )
 
 iface.launch()