Added hermes model

main.py

@@ -1,5 +1,5 @@
 from flask import Flask, request, render_template, jsonify
-from models import llama3
+from models import hermes
 
 app = Flask("AI API")
 
@@ -15,8 +15,20 @@ def test_route():
 def receive_data():
   data = request.get_json()
   print("Prompt:", data["prompt"])
-
-  generated_text = llama3.generate(data["prompt"])
+  
+  messages = []
+  
+  if data["system"]:
+    messages.append({"role": "system", "content": data["system"] })
+  
+  messages.append(
+    {
+      "role": "user",
+      "content": data["prompt"]
+    }
+  )
+  
+  generated_text = hermes.generate(messages)
   
   print("Response:", generated_text)
   

models/hermes.py

@@ -0,0 +1,18 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+model_name = "NousResearch/Hermes-2-Pro-Llama-3-8B"
+
+model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Example messages input
+# messages = [
+# {"role": "system", "content": "You are Hermes 2."},
+# {"role": "user", "content": "Hello, who are you?"}
+#]
+
+def generate(messages):
+  gen_input = tokenizer.apply_chat_template(messages, return_tensors="pt")
+  output_ids = model.generate(**gen_input, num_beams=5, no_repeat_ngram_size=2)
+  
+  return tokenizer.decode(output_ids[0], skip_special_tokens=True)