Update helper/utils.py

helper/utils.py

@@ -6,6 +6,7 @@ import numpy as np
 import pandas as pd
 import PyPDF2
 from openai import OpenAI
+from together import Together
 
 
 # Credit
@@ -168,6 +169,36 @@ def call_gpt(prompt: str, content: str) -> str:
     return response.choices[0].message.content
 
 
+client = Together(api_key=os.environ["TOGETHER_API_KEY"])
+
+def call_llama(prompt: str) -> str:
+    """
+        Send a prompt to the Llama model and return the response.
+        Args:
+            prompt (str): The input prompt to send to the Llama model.
+        Returns:
+            str: The response from the Llama model.
+    """
+
+    # Create a completion request with the prompt
+    response = client.chat.completions.create(
+
+        # Use the Llama-3-8b-chat-hf model
+        model="meta-llama/Llama-3-8b-chat-hf",
+
+        # Define the prompt as a user message
+        messages=[
+            {
+                "role": "user",
+                "content": prompt  # Use the input prompt
+            }
+        ],
+    )
+
+    # Return the content of the first response message
+    return response.choices[0].message.content
+
+
 def quantize_to_kbit(arr: Union[np.ndarray, Any], k: int = 16) -> np.ndarray:
     """Converts an array to a k-bit representation by normalizing and scaling its values.