Update ai_generate.py

ai_generate.py

@@ -4,35 +4,44 @@ import os
 from transformers import pipeline
 from groq import Groq
 
+groq_client = Groq(
+     api_key=os.environ.get("groq_key"),
+)
+
+def generate_groq(text, model):
+    completion = client.chat.completions.create(
+        model = model,
+        messages=[
+            {
+                "role": "user",
+                "content": text
+            },
+            {
+                "role": "assistant",
+                "content": "Please follow the instruction and write about the given topic in approximately the given number of words"
+            }    
+        ],
+        temperature=1,
+        max_tokens=1024,
+        top_p=1,
+        stream=True,
+        stop=None,
+    )
+    response = ""
+    for i, chunk in enumerate(completion):
+        if i != 0:
+            response += chunk.choices[0].delta.content or ""
+    return response
+
 def generate(text, model, api):
     if model == "Llama 3":
-        client = Groq(
-             api_key=os.environ.get("groq_key"),
-        )
-        completion = client.chat.completions.create(
-            model="llama3-8b-8192",
-            messages=[
-                {
-                    "role": "user",
-                    "content": text
-                },
-                {
-                    "role": "assistant",
-                    "content": "Please follow the instruction and write about the given topic in approximately the given number of words"
-                }    
-            ],
-            temperature=1,
-            max_tokens=1024,
-            top_p=1,
-            stream=True,
-            stop=None,
-        )
-
-        response = ""
-        for i, chunk in enumerate(completion):
-            if i != 0:
-                response += chunk.choices[0].delta.content or ""
-        return response
+        return generate_groq(text, "llama3-70b-8192")
+    elif model == "Groq":
+        return generate_groq(text, "llama3-groq-70b-8192-tool-use-preview")
+    elif model == "Mistral":
+        return generate_groq(text, "mixtral-8x7b-32768")    
+    elif model == "Gemma":
+        return generate_groq(text, "gemma2-9b-it")    
     elif model == "OpenAI GPT 3.5":
         client = OpenAI(
             api_key=api,