Jacobsoto's picture
Upload folder using huggingface_hub
cbdc6bc verified
Raw History Blame Contribute Delete
2.76 kB
import os
import sys
from huggingface_hub import InferenceClient
class ChatBot:
def __init__(self, model_name="HelpingAI/HELVETE-3B"):
"""Initialize the chatbot with the specified model."""
# Get API token from environment variable
self.api_key = os.environ.get("HF_TOKEN")
if not self.api_key:
print("Error: HF_TOKEN environment variable not set")
print("Please set your Hugging Face token with: export HF_TOKEN=your_token_here")
sys.exit(1)
self.client = InferenceClient(api_key=self.api_key)
self.model = model_name
self.conversation_history = []
def add_message(self, role, content):
"""Add a message to the conversation history."""
self.conversation_history.append({"role": role, "content": content})
def get_response(self, user_message):
"""Get a response from the AI model."""
# Add user message to history
self.add_message("user", user_message)
try:
# Get completion from the model
completion = self.client.chat.completions.create(
model=self.model,
messages=self.conversation_history,
)
# Extract the assistant's response
assistant_response = completion.choices[0].message.content
# Add assistant response to history
self.add_message("assistant", assistant_response)
return assistant_response
except Exception as e:
return f"Error: {str(e)}"
def chat_loop(self):
"""Run an interactive chat loop."""
print("=" * 50)
print("πŸ€– HELVETE-3B ChatBot")
print("Type 'quit', 'exit', or 'bye' to end the conversation")
print("=" * 50)
while True:
try:
user_input = input("\nYou: ").strip()
if user_input.lower() in ['quit', 'exit', 'bye']:
print("Goodbye! πŸ‘‹")
break
if not user_input:
continue
print("AI: ", end="", flush=True)
response = self.get_response(user_input)
print(response)
except KeyboardInterrupt:
print("\n\nConversation interrupted. Goodbye! πŸ‘‹")
break
except Exception as e:
print(f"\nAn error occurred: {e}")
def main():
"""Main function to run the chatbot."""
# Create and run the chatbot
bot = ChatBot()
bot.chat_loop()
if __name__ == "__main__":
main()