Update app.py

app.py

@@ -3,7 +3,14 @@ import torch
 import threading
 import spaces
 
-from transformers import TextIteratorStreamer
+from transformers import AutoTokenizer, TextIteratorStreamer
+from auto_gptq import AutoGPTQForCausalLM
+
+# Model identifier
+model_id = "xmadai/Mistral-Large-Instruct-2407-xMADai-INT4"
+
+# Load the tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False, trust_remote_code=True)
 
 print("Is CUDA available?", torch.cuda.is_available())
 
@@ -20,6 +27,7 @@ class ModelWrapper:
                 device_map={'': 'cuda:0'},
                 trust_remote_code=True,
             )
+            self.model.eval()
 
         print("Model is on device:", next(self.model.parameters()).device)
 
@@ -48,18 +56,19 @@ class ModelWrapper:
             generated_text += new_text
             yield generated_text
 
+# Instantiate the model wrapper
 model_wrapper = ModelWrapper()
 
+# Create the Gradio interface
 interface = gr.Interface(
     fn=model_wrapper.generate,
     inputs=gr.Textbox(lines=5, label="Input Prompt"),
-    outputs=gr.Textbox(label="Generated Text"),
+    outputs=gr.Textbox(label="Generated Text", lines=10, streaming=True),
     title="Mistral-Large-Instruct-2407 Text Completion",
     description="Enter a prompt and receive a text completion using the Mistral-Large-Instruct-2407 INT4 model.",
     allow_flagging='never',
     live=False,
-    cache_examples=False,
-    streaming=True
+    cache_examples=False
 )
 
 if __name__ == "__main__":