add text gen

main.py

@@ -6,7 +6,9 @@ import re
 
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
+
 import chatbot
+import textgen
 
 
 app = FastAPI()
@@ -31,15 +33,9 @@ class HistoryRequest(BaseModel):
     user: list[str]
     ai: list[str]
 
-@app.post("/generate")
+@app.post("/chatbot")
 async def generate_response(history: HistoryRequest):
-    # try:
-    #     model
-    #     print(12321323)
-    # except: 
-    #     global model
-    #     # model = chatbot.ChatBot()
-
+    print("Chatbot request")
     global model
 
     try:
@@ -63,4 +59,32 @@ async def generate_response(history: HistoryRequest):
     response = model.generate_response(history)
     return response
 
-    
+
+class TextGenInput(BaseModel):
+    user: str
+
+@app.post("/text-gen")
+async def generate_text(input: TextGenInput):
+    print("Generating text request")
+    global model
+    directory = 'models/fine-tuned-gpt2-textgen'
+
+    try:
+        # check if model is already loaded
+        if not isinstance(model, textgen.TextGen):
+            model = textgen.TextGen(
+                GPT2Tokenizer.from_pretrained(directory),
+                GPT2LMHeadModel.from_pretrained(directory),
+                torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            )
+    except NameError:
+        # if model is not defined, load
+        model = textgen.TextGen(
+            GPT2Tokenizer.from_pretrained(directory),
+            GPT2LMHeadModel.from_pretrained(directory),
+            torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        )
+
+    response = model.generate_text(input.user)
+
+    return response

models/fine-tuned-gpt2-textgen/config.json

@@ -0,0 +1,39 @@
+{
+  "_name_or_path": "gpt2",
+  "activation_function": "gelu_new",
+  "architectures": [
+    "GPT2LMHeadModel"
+  ],
+  "attn_pdrop": 0.1,
+  "bos_token_id": 50256,
+  "embd_pdrop": 0.1,
+  "eos_token_id": 50256,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "gpt2",
+  "n_ctx": 1024,
+  "n_embd": 768,
+  "n_head": 12,
+  "n_inner": null,
+  "n_layer": 12,
+  "n_positions": 1024,
+  "reorder_and_upcast_attn": false,
+  "resid_pdrop": 0.1,
+  "scale_attn_by_inverse_layer_idx": false,
+  "scale_attn_weights": true,
+  "summary_activation": null,
+  "summary_first_dropout": 0.1,
+  "summary_proj_to_labels": true,
+  "summary_type": "cls_index",
+  "summary_use_proj": true,
+  "task_specific_params": {
+    "text-generation": {
+      "do_sample": true,
+      "max_length": 50
+    }
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.42.4",
+  "use_cache": true,
+  "vocab_size": 50257
+}

models/fine-tuned-gpt2-textgen/generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 50256,
+  "eos_token_id": 50256,
+  "transformers_version": "4.42.4"
+}

models/fine-tuned-gpt2-textgen/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a712602a48d8498e3836a5fe746ebdbc5aeefa1e7ee3175fee71ae21cad5b8f5
+size 497774208

models/fine-tuned-gpt2-textgen/special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<|endoftext|>",
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

models/fine-tuned-gpt2-textgen/tokenizer_config.json

@@ -0,0 +1,22 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "50256": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

textgen.py

@@ -0,0 +1,30 @@
+import torch
+
+class TextGen:
+    def __init__(self,tokenizer,model,device):
+        self.tokenizer = tokenizer
+        self.model = model
+        self.device = device
+        self.model.to(self.device)
+
+
+    def generate_text(self, user_input):
+        inputs = self.tokenizer.encode(user_input, return_tensors="pt").to(self.device)
+
+        # generate text
+        attention_mask = torch.ones(inputs.shape, device=self.device)
+        output = self.model.generate(
+            inputs,
+            attention_mask=attention_mask,
+            num_return_sequences=1,
+            max_length=50,
+            max_new_tokens=100,
+            temperature=0.5,
+            repetition_penalty=1.2,
+            pad_token_id=self.tokenizer.eos_token_id,
+            )
+
+        generated_text = self.tokenizer.decode(output[0], skip_special_tokens=True)
+
+
+        return generated_text