vimedllm / notebook /An /new /chatbot.py
VuvanAn's picture
Upload folder using huggingface_hub
cc37925 verified
# Class Chatbot
import camel
from camel.models import ModelFactory
from camel.types import ModelType, ModelPlatformType
from camel.configs import MistralConfig, OllamaConfig, ChatGPTConfig
import time
import os
class Chatbot:
def __init__(self, model_name: str = "gpt", max_token: int = 1, temperature:float = 0.3):
self.model_name = model_name.lower()
self.max_token = max_token
self.temperature = temperature
self.model = self._get_model()
def _get_model(self):
if "gpt" == self.model_name:
return ModelFactory.create(
model_platform=ModelPlatformType.OPENAI,
model_type=ModelType.GPT_4O_MINI,
model_config_dict=ChatGPTConfig(temperature=0.3, max_tokens=self.max_token).as_dict(),
)
elif "mistral" == self.model_name:
return ModelFactory.create(
model_platform=ModelPlatformType.MISTRAL,
model_type=ModelType.MISTRAL_MEDIUM_3,
model_config_dict=MistralConfig(temperature=0.3, max_tokens=self.max_token).as_dict(),
)
else:
return ModelFactory.create(
model_platform=ModelPlatformType.OLLAMA,
model_type=self.model_name,
model_config_dict=OllamaConfig(temperature=0.3, max_tokens=self.max_token).as_dict(),
)
def chat(self, user_input: str, system:str = "You are a helpful assistant.", suppress_error: bool = False) -> str:
messages = [
{"role": "system", "content": system},
{"role": "user", "content": user_input},
]
while True: # Retry loop
try:
response = self.model.run(messages)
return response.choices[0].message.content
except Exception as e:
if suppress_error == False:
print(f"Error: {e}")
time.sleep(2)
# # Example usage
if __name__ == "__main__":
chatbot_mistral = Chatbot(model_name="mistral")
print(chatbot_mistral.chat("hi"))
# def ask_mistral(sys: str, user: str) -> str:
# response = chatbot_mistral.chat(user_input=user, system=sys)
# return response
# chatbot_ollama = Chatbot(model_name="gemma3:12b")
# def ask_ollama(sys: str, user: str) -> str:
# response = chatbot_ollama.chat(user_input=user, system=sys)
# return response
# # Initialize chatbot with OpenAI model
# chatbot = Chatbot(model_name="mistral")
# user_input = "How parasites affect the heart, especially heart failure? Which parasites cause these problems? And how can they do?"
# response = chatbot.chat(user_input)
# print("Chatbot response:", response)