chat.py impl

config.py

@@ -4,8 +4,8 @@ from dotenv import load_dotenv
 # Load from .env file. Store your HF token in the .env file.
 load_dotenv()
 
-
-BASE_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+BASE_MODEL = "HuggingFaceH4/zephyr-7b-beta"
+# BASE_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 # Other options:
 # MODEL = "meta-llama/Llama-2-7b-chat-hf"
 # MODEL = "openlm-research/open_llama_3b"

src/chat.py

@@ -37,7 +37,11 @@ class SchoolChatbot:
              User: {user_input}
              Assistant:"
         """
-        pass
+        return (
+            "<|system|>You are a helpful assistant that specializes in Boston public school enrollment.<|end|>\n"
+            f"<|user|>{user_input}<|end|>\n"
+            "<|assistant|>"
+        )
         
     def get_response(self, user_input):
         """
@@ -58,4 +62,12 @@ class SchoolChatbot:
         - Use self.format_prompt() to format the user's input
         - Use self.client to generate responses
         """
-        pass
+        prompt = self.format_prompt(user_input)
+        response = self.client.text_generation(
+            prompt,
+            max_new_tokens=512,
+            temperature=0.7,
+            top_p=0.9,
+            stop_sequences=["<|end|>"]
+        )
+        return response.strip()