Update main.py

main.py

@@ -1,11 +1,7 @@
-
 from flask import Flask, request, jsonify
 from huggingface_hub import InferenceClient
 
-
-client = InferenceClient(
-    "mistralai/Mistral-7B-Instruct-v0.1"
-)
+client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.1")
 
 app = Flask(__name__)
 
@@ -17,13 +13,18 @@ with open(file_path, "r") as file:
 def home():
     return jsonify({"message": "Welcome to the Recommendation API!"})
 
-
 def format_prompt(message):
     prompt = "<s>"
     prompt += f"[INST] {message} [/INST]"
     prompt += "</s>"
     return prompt
 
+def generate_output(stream):
+    output = ""
+    for response in stream:
+        output += response.token.text
+        yield output
+
 @app.route('/get_course', methods=['POST'])
 def recommend():
     temperature = 0.9
@@ -54,26 +55,23 @@ def recommend():
     {{"course1:course_name, course2:course_name, course3:course_name,...}}
     """
     formatted_prompt = format_prompt(prompt)
-    print(formatted_prompt)
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
 
-    for response in stream:
-        output += response.token.text
-        yield output
-    return output
+    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+    return jsonify({"ans": list(generate_output(stream))})
 
 @app.route('/get_mentor', methods=['POST'])
 def mentor():
-    temperature=0.9
-    max_new_tokens=256
-    top_p=0.95
-    repetition_penalty=1.0
+    temperature = 0.9
+    max_new_tokens = 256
+    top_p = 0.95
+    repetition_penalty = 1.0
+
     content = request.json
     user_degree = content.get('degree')
     user_stream = content.get('stream')
     user_semester = content.get('semester')
     courses = content.get('courses')
+
     temperature = float(temperature)
     if temperature < 1e-2:
         temperature = 1e-2
@@ -89,12 +87,10 @@ def mentor():
     )
     prompt = f""" prompt:
     You need to act like as recommendataion engine for mentor recommendation for student based on below details also the list of mentors with their experience is attached.
-
     Degree: {user_degree}
     Stream: {user_stream}
     Current Semester: {user_semester}
     courses opted:{courses}
-
     Mentor list= {mentors_data}
     Based on above details recommend the mentor that realtes to above details 
     Note: Output should be valid json format in below format:
@@ -103,13 +99,7 @@ def mentor():
     formatted_prompt = format_prompt(prompt)
 
     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
-
-    for response in stream:
-        output += response.token.text
-        yield output
-    return jsonify({"ans":output})
-
+    return jsonify({"ans": list(generate_output(stream))})
 
 if __name__ == '__main__':
-    app.run(debug=True)
+    app.run(debug=True)