Update llama2.py

llama2.py

@@ -2,6 +2,7 @@ import os
 import json
 import requests
 import sseclient
+import google.generativeai as palm_api
 
 from pingpong import PingPong
 from pingpong.pingpong import PPManager
@@ -9,36 +10,14 @@ from pingpong.pingpong import PromptFmt
 from pingpong.pingpong import UIFmt
 from pingpong.gradio import GradioChatUIFmt
 
-class LLaMA2ChatPromptFmt(PromptFmt):
-    @classmethod
-    def ctx(cls, context):
-        if context is None or context == "":
-            return ""
-        else:
-            return f"""<<SYS>>
-{context}
-<</SYS>>
-"""
 
-    @classmethod
-    def prompt(cls, pingpong, truncate_size):
-        ping = pingpong.ping[:truncate_size]
-        pong = "" if pingpong.pong is None else pingpong.pong[:truncate_size]
-        return f"""[INST] {ping} [/INST] {pong}"""
+palm_api_token = os.getenv("PALM_API_TOKEN")
+if palm_api_token is None:
+    raise ValueError("PaLM API Token is not set")
+else:
+    palm_api.configure(api_key=palm_api_token)
 
-class LLaMA2ChatPPManager(PPManager):
-    def build_prompts(self, from_idx: int=0, to_idx: int=-1, fmt: PromptFmt=LLaMA2ChatPromptFmt, truncate_size: int=None):
-        if to_idx == -1 or to_idx >= len(self.pingpongs):
-            to_idx = len(self.pingpongs)
-
-        results = fmt.ctx(self.ctx)
-
-        for idx, pingpong in enumerate(self.pingpongs[from_idx:to_idx]):
-            results += fmt.prompt(pingpong, truncate_size=truncate_size)
-
-        return results
-
-class GradioLLaMA2ChatPPManager(LLaMA2ChatPPManager):
+class GradioPaLMChatPPManager(PPManager):
     def build_uis(self, from_idx: int=0, to_idx: int=-1, fmt: UIFmt=GradioChatUIFmt):
         if to_idx == -1 or to_idx >= len(self.pingpongs):
             to_idx = len(self.pingpongs)
@@ -48,86 +27,39 @@ class GradioLLaMA2ChatPPManager(LLaMA2ChatPPManager):
         for pingpong in self.pingpongs[from_idx:to_idx]:
             results.append(fmt.ui(pingpong))
 
-        return results
+        return results    
 
-async def gen_text(
-    prompt, 
-    hf_model='meta-llama/Llama-2-70b-chat-hf', 
-    hf_token=None, 
+def gen_text(
+    prompt,
+    palm,
     parameters=None
 ):
-  if hf_token is None:
-    raise ValueError("Hugging Face Token is not set")
-
-  if parameters is None:
-    parameters = {
-        'max_new_tokens': 512,
-        'do_sample': True,
-        'return_full_text': False,
-        'temperature': 1.0,
-        'top_k': 50,
-        # 'top_p': 1.0,
-        'repetition_penalty': 1.2
-    }
-
-  url = f'https://api-inference.huggingface.co/models/{hf_model}'
-  headers={
-      'Authorization': f'Bearer {hf_token}',
-      'Content-type': 'application/json'
-  }
-  data = {
-      'inputs': prompt,
-      'stream': True,
-      'options': {
-          'use_cache': False,
-      },
-      'parameters': parameters
-  }
-
-  r = requests.post(
-      url,
-      headers=headers,
-      data=json.dumps(data),
-      stream=True
-  )
-
-  client = sseclient.SSEClient(r)
-  for event in client.events():
-    yield json.loads(event.data)['token']['text']
-
-def gen_text_none_stream(
-    prompt, 
-    hf_model='meta-llama/Llama-2-70b-chat-hf', 
-    hf_token=None, 
-):
-    parameters = {
-        'max_new_tokens': 64,
-        'do_sample': True,
-        'return_full_text': False,
-        'temperature': 0.7,
-        'top_k': 10,
-        # 'top_p': 1.0,
-        'repetition_penalty': 1.2
-    }
-
-    url = f'https://api-inference.huggingface.co/models/{hf_model}'
-    headers={
-        'Authorization': f'Bearer {hf_token}',
-        'Content-type': 'application/json'
-    }
-    data = {
-        'inputs': prompt,
-        'stream': False,
-        'options': {
-            'use_cache': False,
-        },
-        'parameters': parameters
-    }
-
-    r = requests.post(
-        url,
-        headers=headers,
-        data=json.dumps(data),
-    )
-
-    return json.loads(r.text)[0]["generated_text"]
+    if parameters is None:
+        model = 'models/text-bison-001'
+        temperature = 0.7
+        candidate_count = 1
+        top_k = 40
+        top_p = 0.95
+    
+        parameters = {
+            'model': model,
+            'temperature': temperature,
+            'candidate_count': candidate_count,
+            'top_k': top_k,
+            'top_p': top_p,
+        }
+
+    if palm is None:
+        response = palm_api.chat(messages=[prompt])
+    else:
+        response = palm.reply(prompt)
+
+    if len(response.filters) > 0 and \
+        response.filters[0]['reason'] == 2:
+        response_txt = "your request is blocked for some reasons"
+    
+    else:
+        response_txt = response.last
+    
+    return response, response_txt
+