Update helpers/foundation_models.py

helpers/foundation_models.py

@@ -147,3 +147,50 @@ def llama2_7b_ysa(prompt: str) -> str:
     response: str = output[0]['generated_text']
 
     return response
+
+
+def llama2_7b_brk_letters(prompt: str) -> str:
+    """
+    Queries the LLaMA 2 7B model hosted on a specific Hugging Face endpoint with a given prompt,
+    and returns the generated text as a response.
+
+    Args:
+        prompt (str): The input text prompt to be sent to the LLaMA 2 7B model for generating text.
+
+    Returns:
+        str: The text generated by the LLaMA 2 7B model in response to the input prompt.
+    """
+    
+    # Endpoint URL of the LLaMA 2 7B model hosted on Hugging Face.
+    API_URL = "https://hd1rl4q31aom5qwc.us-east-1.aws.endpoints.huggingface.cloud"
+    
+    # Headers to include in the HTTP request. Specifies the expected format of the response and request.
+    headers = {
+        "Accept": "application/json",
+        "Content-Type": "application/json"
+    }
+    
+    def query_llama2_7b_brk_letters(payload: dict) -> dict:
+        """
+        Sends a POST request to the LLaMA 2 7B API endpoint with a given payload.
+
+        Args:
+            payload (dict): The data to be sent in the POST request, including the input prompt
+            and any parameters for the model.
+
+        Returns:
+            dict: The JSON response from the API, parsed into a dictionary.
+        """
+        response = requests.post(API_URL, headers=headers, json=payload)
+        return response.json()
+    
+    # The payload for the POST request, including the prompt and any model parameters.
+    output = query_llama2_7b_brk_letters({
+        "inputs": prompt,
+        "parameters": {}
+    })
+
+    # Extracts the generated text from the API's response.
+    response = output[0]['generated_text']
+
+    return response