add second version

app.py

@@ -1,3 +1,4 @@
+import openai
 import time
 import yaml
 import gradio as gr
@@ -8,11 +9,14 @@ from chatbot import ChatBot
 
 # MARKDOWN
 MARKDOWN = """
-# Coffee-Gym Feedback Model
-Welcome to the COFFEE-GYM demo page! This page will guide you through using our comprehensive RL environment for training models that provide feedback on code editing.
+# Coffee-Gym
+Welcome to the COFFEE-GYM demo page! 
+
+## DS-Coder-7B-PPO-CoffeeEval
+You can try using the COFFEE-GYM feedback model (DS-Coder-7B-PPO-CoffeeEval) to get feedback on your code.
 
 ## Prompt template
-To use the COFFEE-GYM feedback model, you can follow the prompt template below:
+To use the DS-Coder-7B-PPO-CoffeeEvall, you can follow the prompt template below:
 
 ~~~json
 Problem Description:
@@ -33,9 +37,20 @@ Feel free to explore and make the most out of COFFEE-GYM!
 """
 
 
-def main():
+
+
+# get arguments from command line
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_name", type=str, default="Team-Coffee-Gym/DS-Coder-7B-PPO-CoffeeEval", help="Model name to use for the chatbot.")
+    parser.add_argument("--url", type=str, default="https://editor.jp.ngrok.io/v1", help="URL to use for the chatbot.")
+    args = parser.parse_args()  
+    return args
+
+
+def main(args):
     
-    chatbot = ChatBot()
+    chatbot = ChatBot(model_name=args.model_name, url=args.url)
     
     with gr.Blocks() as app:
         ##### Playground #####
@@ -70,4 +85,5 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    args = get_args()
+    main(args)

chatbot.py

@@ -1,19 +1,36 @@
+import openai
 import time
 import json
 import yaml
 from typing import Union
 import os
-
-
-from transformers import pipeline
+from langchain_openai import OpenAI
+from transformers import AutoTokenizer
+from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
+from langchain_community.chat_models import ChatOpenAI
+from langchain.callbacks import get_openai_callback
 
 
 class ChatBot(object):
-    def __init__(self):
-        self.model = pipe = pipeline("text-generation", model="Team-Coffee-Gym/CoffeeGym")
+    def __init__(self, model_name, url):
+        self.model_name = model_name
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.url = url
+        self.model = self._set_model()
         self.chat_history = []
         self.history = self._set_initial_history()
             
+
+    def _set_model(self):
+        return OpenAI(
+            model_name=self.model_name,
+            openai_api_base = f"{self.url}",
+            openai_api_key="EMPTY",
+            temperature = 0.1,
+            top_p = 0.95,
+            max_tokens = 1024,
+            max_retries=3
+        )
         
     def _set_initial_history(self):
         return ["You are an exceptionally intelligent coding assistant developed by DLI lab that consistently delivers accurate and reliable responses to user instructions. If somebody asks you who are you, answer as 'AI programming assistant based on DLI Lab'.\n\n"]
@@ -50,14 +67,14 @@ class ChatBot(object):
  
         self.chat_history = chat_history
         
-        model_input = self.set_model_input(input_text)
-        response = self.model(model_input)
+        model_input_list = self.set_model_input(input_text)
+        model_input = self.tokenizer.apply_chat_template(model_input_list, tokenize=False, add_generation_prompt=True)
+        response = self.model.invoke(model_input)
         
         if response is not None:
             self.history.append(response)
             self.chat_history = self.chat_history + [(input_text, response)]
         
-        self.log_chat()
         return self.chat_history
     
     
@@ -66,7 +83,8 @@ class ChatBot(object):
         self.chat_history = chat_history[:-1]
         self.history = self.history[:-2]
         
-        model_input = self.set_model_input(None)
+        model_input_list = self.set_model_input(None)
+        model_input = self.tokenizer.apply_chat_template(model_input_list, tokenize=False, add_generation_prompt=True)
         response = self.model.invoke(model_input)
         
         if response is not None:

requirements.txt

@@ -1,4 +1,5 @@
 huggingface_hub==0.23.0
-transformers==4.41.2
-gradio==4.37.2
-gradio-client==1.0.2
+transformers
+gradio
+langchain
+langchain_openai