Update app.py

app.py

@@ -3,11 +3,17 @@ import requests
 import os
 import logging
 from datetime import datetime
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import TextStreamer
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-device = "cuda"
+# Load the model and tokenizer
+model_name = "mixtral/instruct-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
 
-@spaces.GPU()
 def stream_chat(
     message: str,
     history: list,
@@ -18,9 +24,6 @@ def stream_chat(
     top_k: int = 20,
     penalty: float = 1.2,
 ):
-    print(f'message: {message}')
-    print(f'history: {history}')
-
     conversation = [
         {"role": "system", "content": system_prompt}
     ]
@@ -32,9 +35,9 @@ def stream_chat(
 
     conversation.append({"role": "user", "content": message})
 
-    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(model.device)
+    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(device)
 
-    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+    streamer = TextStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
 
     generate_kwargs = dict(
         input_ids=input_ids,
@@ -46,6 +49,10 @@ def stream_chat(
         eos_token_id=[128001, 128008, 128009],
         streamer=streamer,
     )
+
+    output = model.generate(**generate_kwargs)
+    return tokenizer.decode(output[0], skip_special_tokens=True)
+
 app = Flask(__name__)
 
 # Configure logging