Update modules/chatbot.py

modules/chatbot.py

@@ -3,14 +3,49 @@ import torch
 
 class MultilingualChatbot:
     def __init__(self):
-        self.model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
-        self.tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+        self.models = {
+            'en': GPT2LMHeadModel.from_pretrained("microsoft/DialoGPT-medium"),
+            'es': GPT2LMHeadModel.from_pretrained("DeepESP/gpt2-spanish"),
+            'fr': GPT2LMHeadModel.from_pretrained("asi/gpt-fr-cased-small")
+        }
+        self.tokenizers = {
+            'en': GPT2Tokenizer.from_pretrained("microsoft/DialoGPT-medium"),
+            'es': GPT2Tokenizer.from_pretrained("DeepESP/gpt2-spanish"),
+            'fr': GPT2Tokenizer.from_pretrained("asi/gpt-fr-cased-small")
+        }
+        for tokenizer in self.tokenizers.values():
+            tokenizer.pad_token = tokenizer.eos_token
 
     def generate_response(self, prompt, src_lang):
-        self.tokenizer.src_lang = src_lang
-        encoded_input = self.tokenizer(prompt, return_tensors="pt")
-        generated_tokens = self.model.generate(**encoded_input, max_length=100)
-        return self.tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+        # Default to English if the language is not supported
+        model = self.models.get(src_lang, self.models['en'])
+        tokenizer = self.tokenizers.get(src_lang, self.tokenizers['en'])
+        
+        input_ids = tokenizer.encode(prompt + tokenizer.eos_token, return_tensors='pt')
+        
+        # Move input to the same device as the model
+        input_ids = input_ids.to(model.device)
+        
+        chat_history_ids = model.generate(
+            input_ids,
+            max_length=1000,
+            pad_token_id=tokenizer.eos_token_id,
+            no_repeat_ngram_size=3,
+            do_sample=True,
+            top_k=50,
+            top_p=0.95,
+            temperature=0.7,
+            num_return_sequences=1,
+            length_penalty=1.0,
+            repetition_penalty=1.2
+        )
+        return tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
+
+def initialize_chatbot():
+    return MultilingualChatbot()
+
+def get_chatbot_response(chatbot, prompt, src_lang):
+    return chatbot.generate_response(prompt, src_lang)
 
 def initialize_chatbot():
     return MultilingualChatbot()