Update ai_generate.py

ai_generate.py

@@ -2,19 +2,35 @@ import torch
 from openai import OpenAI
 import os
 from transformers import pipeline
-
-# pipes = {
-#     'GPT-Neo': pipeline("text-generation", model="EleutherAI/gpt-neo-2.7B"),
-#     'Llama 3': pipeline("text-generation", model="meta-llama/Meta-Llama-3-8B")
-# }
+from groq import Groq
 
 def generate(text, model, api):
-    if model == "GPT-Neo":
-        response = pipeline("text-generation", model="EleutherAI/gpt-neo-2.7B")(text)
-        print(response)
-        return response[0]
-    elif model == "Llama 3":
-        response = pipeline("text-generation", model="meta-llama/Meta-Llama-3-8B")(text)
+    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,
+        )
+        
+        for chunk in completion:
+            print(chunk.choices[0].delta.content or "", end="")
+
         print(response)
         return response[0] 
     elif model == "OpenAI GPT 3.5":