Update main.py

main.py

@@ -3,9 +3,10 @@ from flask import Flask, request, jsonify
 
 device = "cuda" # the device to load the model onto
 
+from ctransformers import AutoModelForCausalLM
+
+llm = AutoModelForCausalLM.from_pretrained("TheBloke/Llama-2-7b-Chat-GGUF", model_file="llama-2-7b-chat.q4_K_M.gguf", model_type="llama", gpu_layers=0)
 
-model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
-tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
 
 
 
@@ -15,8 +16,7 @@ def recommendation():
     user_degree = content.get('degree')
     user_stream = content.get('stream')
     user_semester = content.get('semester')
-    messages = [
-    {"role": "user", "content": f"""
+    prompt = """
     You need to act like as recommendataion engine for course recommendation based on below details.
 
     Degree: {user_degree}
@@ -28,18 +28,11 @@ def recommendation():
     Note: Output should bevalid json format in below format:
     {{"course1:ABC,course2:DEF,course3:XYZ,...}}
     
-    """},
-    
-    ]
-    
-    encodeds = tokenizer.apply_chat_template(messages, 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)
-    return jsonify({"res":decoded[0]})
+    """
+    suffix="[/INST]"
+    prefix="[INST] <<SYS>> You are a helpful assistant <</SYS>>"
+    prompt = f"{prefix}{user.replace('{prompt}', prompt)}{suffix}"
+    return jsonify({"ans":llm(prompt)})
 
 if __name__ == '__main__':
     app.run(debug=True)