Update app.py

app.py

@@ -16,7 +16,45 @@ app = FastAPI(root_path="/api/v1")
 
 # Load the model and tokenizer
 model_name_or_path = "TheBloke/Wizard-Vicuna-7B-Uncensored-GPTQ"
+mistral_model="mistralai/Mistral-7B-Instruct-v0.2"
 conversations = {}
+device = "cuda" # the device to load the model onto
+
+def mistral_model():
+        """
+    Loads the GPT-3.5 model and tokenizer.
+    
+    Returns:
+        tuple: A tuple containing the loaded model and tokenizer.
+    """
+    if torch.cuda.is_available():
+        print("CUDA is available. GPU will be used.")
+    else:
+        print("CUDA is not available. CPU will be used.")
+   
+    model = AutoModelForCausalLM.from_pretrained(mistral_chat)
+
+    tokenizer = AutoTokenizer.from_pretrained(mistral_chat)
+
+    return model,tokenizer
+    
+model, tokenizer = mistral_model()
+
+def mistral_generated_response(msg_prompt, persona_desc_prompt):
+    user_prompt = f'{msg_prompt} [/INST]'
+    persona_prompt = f'{persona_desc_prompt} [/INST]'
+    prompt_template = f'''### [INST] Instruction:{persona_prompt} [INST] {user_prompt}'''
+    
+    encodeds = tokenizer.apply_chat_template(prompt_template, return_tensors="pt")
+    
+    model_inputs = encodeds.to(device)
+    model.to(device)
+    generated_ids = model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
+    decoded = tokenizer.batch_decode(generated_ids)
+    response = (decoded[0])
+    return response
+
+
 
 def load_model_norm():
     """
@@ -35,7 +73,7 @@ def load_model_norm():
 
     return model, tokenizer 
 
-model, tokenizer = load_model_norm()
+#model, tokenizer = load_model_norm()
 
 
 def generate_response(msg_prompt: str) -> dict:
@@ -192,3 +230,16 @@ async def get_response(thread_id: int):
     response = thread['responses'][-1]
 
     return {'response': response}
+
+@app.post("/mistral_chat")
+async def mistral_chat(prompt: dict):
+    try:
+        msg_prompt = prompt.get("msg_prompt")
+        persona_desc_prompt = prompt.get("persona_desc_prompt")
+        if not msg_prompt or not persona_desc_prompt:
+            return {"error": "msg_prompt and persona_desc_prompt are required fields."}
+        
+        response = mistral_generated_response(msg_prompt, persona_desc_prompt)
+        return {"response": response, "prompt": {"msg_prompt": msg_prompt, "persona_desc_prompt": persona_desc_prompt}}
+    except Exception as e:
+        return {"error": str(e)}