update

chatbot.py

@@ -4,27 +4,30 @@ from datasets import load_dataset
 import pandas as pd
 import re
 
-
 class ChatBot:
-    def __init__(self):
-        self.directory = 'models/fine-tuned-gpt2'
-        self.tokenizer = GPT2Tokenizer.from_pretrained(self.directory)
-        self.model = GPT2LMHeadModel.from_pretrained(self.directory)
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    def __init__(self,dir,tokenizer,model,device):
+        self.directory = dir
+        self.tokenizer = tokenizer
+        self.model = model
+        self.device = device
         self.model.to(self.device)
 
     def generate_response(self, history):
         combined_prompt = ""
+
         # self.tokenizer.eos_token_id = '<|endoftext|>'
         if len(history.user) > 7:
             history.user = history.user[-7:]
             history.ai = history.ai[-6:] 
+
         # Iterate over user and AI messages
         for user_message, ai_message in zip(history.user, history.ai):
             combined_prompt += f"<user> {user_message}{self.tokenizer.eos_token_id}<AI> {ai_message}{self.tokenizer.eos_token_id}"
+       
         # Include the last user message in the prompt for response generation
         if history.user:
             combined_prompt += f"<user> {history.user[-1]}{self.tokenizer.eos_token_id}<AI>"
+        
         # Tokenize and generate response
         inputs = self.tokenizer.encode(combined_prompt, return_tensors="pt").to(self.device)
         attention_mask = torch.ones(inputs.shape, device=self.device)

main.py

@@ -1,3 +1,9 @@
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
+import torch
+from datasets import load_dataset
+import pandas as pd
+import re
+
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 import chatbot
@@ -5,7 +11,6 @@ import chatbot
 
 app = FastAPI()
 
-
 # Add CORS middleware to allow any origin
 app.add_middleware(
     CORSMiddleware,
@@ -28,13 +33,32 @@ class HistoryRequest(BaseModel):
 
 @app.post("/generate")
 async def generate_response(history: HistoryRequest):
-    try:
-        model
-    except: 
-        model = chatbot.ChatBot()
+    # try:
+    #     model
+    #     print(12321323)
+    # except: 
+    #     global model
+    #     # model = chatbot.ChatBot()
 
-    if type(model) != type(chatbot.ChatBot()):
-        model = chatbot.ChatBot()
+    global model
+
+    try:
+        # check if model is already loaded
+        if not isinstance(model, chatbot.ChatBot):
+            model = chatbot.ChatBot(
+                'models/fine-tuned-gpt2',
+                GPT2Tokenizer.from_pretrained('models/fine-tuned-gpt2'),
+                GPT2LMHeadModel.from_pretrained('models/fine-tuned-gpt2'),
+                torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            )
+    except NameError:
+        # if model is not defined, load
+        model = chatbot.ChatBot(
+            'models/fine-tuned-gpt2',
+            GPT2Tokenizer.from_pretrained('models/fine-tuned-gpt2'),
+            GPT2LMHeadModel.from_pretrained('models/fine-tuned-gpt2'),
+            torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        )
     
     response = model.generate_response(history)
     return response