Upload 4 files

Inference.py

@@ -0,0 +1,46 @@
+import os, time, torch, warnings
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+
+class Inference():
+
+    def __init__(self, silent=False) -> None:
+        start_time = time.perf_counter()
+        self.tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
+        self.model = GPT2LMHeadModel.from_pretrained(self.local_file_path("SaveState"))
+        self.model.eval()
+        if not silent:
+            print(f"Model Loading Took {time.perf_counter()-start_time} Seconds")
+
+    def local_file_path(self, path):
+        return os.path.join(os.path.dirname(os.path.abspath(__file__)), path)
+
+    def generate(self, prompt, max_length=2000, temperature=0.5, do_sample=True, stop_token=None, callback=None, silent=True):
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            start_time = time.perf_counter()
+            input_ids = self.tokenizer.encode(prompt, return_tensors='pt')
+            generated_text = input_ids
+            while generated_text.shape[1] < max_length:
+                length = min(50, max_length - generated_text.shape[1])
+                with torch.no_grad():
+                    outputs = self.model.generate(input_ids, max_length=length, temperature=temperature, do_sample=do_sample, pad_token_id=self.tokenizer.eos_token_id)
+                new_tokens = outputs[0][-length:]
+                if callback is not None:
+                    for token in new_tokens:
+                        callback(self.tokenizer.decode([token]))
+                generated_text = torch.cat((generated_text, new_tokens.unsqueeze(0)), dim=-1)
+                input_ids = new_tokens.unsqueeze(0)
+                if stop_token is not None and stop_token in self.tokenizer.decode(generated_text[0]):
+                    break
+            if not silent:
+                print(f"Model Loading Took {time.perf_counter()-start_time} Seconds")
+            return self.tokenizer.decode(generated_text[0], skip_special_tokens=True)
+
+Inference = Inference()
+
+def spec(stre):
+    print(stre, end="")
+
+if __name__=="__main__":
+    while True:
+        print(Inference.generate(input(">>> "), max_length=100, temperature=0.8, silent=True))

TrainData.txt

@@ -0,0 +1,118 @@
+Scene: Sunshine
+Person1: Ahh beautiful sunshine
+Person2: I love the way it bounces off your - beautiful face
+Person1: You were going to say something else there
+Person2: I resisted the bald joke because it's still quite early in the show but now you've opened the door so here it goes, you have no hair
+Person1: I know I don't have any hair, but no we should not fight. It is our anniversary night.
+Person2: It is our anniversary
+Person1: Yes
+Person2: And we always come to this picnic spot for our anniversary in the sunshine
+Person1: Now come, let us eat some bread and cheese.
+Person2: For we are...
+Both: French
+Person1: Thank You
+Person2: To us
+Person1: To us
+Both: Clink
+Person2: So how's your affair going
+Person1: ha ha
+Person2: I know. *Spits out wine*
+Person1: What did you just do?
+Person2: I kept the wine in my mouth the whole time because I was going to dramatically reveal that it was poisoned. You're going to die in the sunshine for what you did I know about Jean Claude.
+Person1: Jean Claude!? Jean Claude and I are just buddies!
+Person2: Buddies? More like sex buddies! I turn the corner of bakers street and there you two were raw dogging it in the streets!
+Person1: No we weren't! We were just having good times together having a play roll!
+Person2: A play roll!? You think that I am that naive?
+Person1: Oh-
+Person2: Wow this is a very slow acting poison
+Person3: OH MY GOODNESS!! THAT MAN IS BALD!
+**End**
+Scene: Territory
+Person1: Easy there, you're about to step on my territory. You actually have.
+Person2: Well I see I've made some kind of mistake
+Person1: Well all you gotta do is just wander back over there and we ain't got no problems.
+Person2: I just gotta wander back over there and we ain't gonna have a problem?
+Person1: Oh what, are you repeating everything I'm saying? What is there an echo in here?
+Person2: I JUST WANNA BE REAL CLEAR ON WHATS I GOTS TO DO! I JUST HAVE TO WALK OVER THERE AND WE AIN'T GONNA HAVE NO MORE PROBLEMS IS THAT WHAT YA SAID!?
+Person1: I ain't saying we're gonna not have no problems I mean we still is gonna have taxation and relationship troubles. I'm saying the problem of me having a gun pointed on you is gonna diminish awful quick like
+Person2: So you were being over generalized before when you said we ain't gonna have no problems you meant this new specific problem will be gone but we'll still have problems such as inflation.
+Person1: I apoligize for not being hyper-specific with my initial sentence. I will try to do better in the future. Now please move exactly two meters to your left...
+Person2: I UNDERSTAND YOUR MISTAKE THAT RHETORIC CAN BE A TRICKY MISTRESS I WANT TO UNDERSTAND PRECISELY WHAT WILL OCCUR IF I'M TO STEP BACK OVER THERE?
+Person3: Hey if you do too far you'll be on my property and I don't want there to be any issues because thats my-
+Person2: DO YOU UNDERSTAND WHAT HE'S SAYIN?
+Person1: I dont give a f*** about what he's saying.
+Person2: We have no idea what your sayin you made this a 3 body problem here which is an unsolvable mathematical issue.
+Person3: I'm just saying that this is my turf right here okay-
+Person2: only you are on my best I do not see why
+Person1: When you say thats your turf do you only own the turf or do you own the land underneath it. If we don't go on the turf are we still allowed over there or is it a property line? Be more specific!
+**End**
+Scene: Be More Specific
+Person1: Just be more specific about it.
+Person2: Okay okay so... like...
+Person1: Cmon!
+Person2: Im leaving for work, right? I hear a knock at the door and its the delivery driver and I say "Oh hello!" and he says "Hello!" and he is delivering a parcel and he says its for me and he hands me a stylus to sign for it and my hand slips and the stylus goes through the device into his heart and now he's dead.
+Person1: Okay you're underarrest.
+**End**
+Scene: Hunting
+Person1: I've been hunting wild and exotic animals for about two weeks now and I realize I really should have prepared a lot more, as you can see I've lost me legs and me arms. Turns out crocodiles, even if they look like they're smiling, aren't actually your friends. Thanks for coming to the Ted Talk. Don't f*** with nature 'cause nature f***s back. Now what I wanna recommend for anyone who wants to do what I did is do your research. Do your reading and start with something small like a doll house. Very unlikely a doll house is going to give you such grievous bodily injuries. Nexy slide!
+Person2: *Snaps Fingers*
+Person1: There's Jim, he's my assistant. Since the incident I can't click anymore so we made a click activated powerpoint presentation. As you can see this is a photo of the incident, actually about 2 seconds before, and I am now waving at the crocodile in the water.
+Person2: *Snaps Fingers*
+Person1: Now I'm really in water. As you can see one of the arms is gone. Why did I not run away then you might ask? I've always been a stubborn man.
+**End**
+Scene: It's Okay
+Person1: Hey it's okay
+Person2:I just um... I just need this time.
+Person1: You want us to leave?
+Person2: Yeah
+Person1: *Leaves*
+Person3: *Leaves*
+Person2: No wait don't leave.
+Person3: Listen mate it will be okay
+Person2: I don't want to be in a solo scene
+Person1: That was my bad mate I really thought someone might come in at some point
+Person2: I just... sometimes it's funny just to let someone off stage
+Person3: Listen it's alright if you can't think of anything
+Person1: I know its your first time doing improve Person2
+Person2: Its just I've run out of ideas
+Person1: Well I am sure you can think of something
+Person2: You just have to be funny every moment and it's just f***ing exhausting. Plus we are running out of accents we've done Australian, German, French, British...
+Person3: Don't do the offensive ones. Leave those in the nineties
+Person2: I don't know mate I think I should just quit.
+Person1: We don't want you to go... We just gotta do a couple more games in this half then have a break and a lemonade and just come back for the second half, alright?
+**End**
+Scene: Digging
+Person1: That's good Laddy, Listen we need to get these bodies in here by the time the police come alright
+Person2: and these
+Person1: and th- f***in ell how many people?
+Person2: sixpence word 6 pennies
+Person1: My goodness. We've robbed this bank okay and it's gone a bit wrong there's some- we've killed a few people all right this is...
+Person3: violence criminel
+Person1: I know this is what we are. we're violent criminals. We are violent criminals! Person3 are you alright?
+Person3: Volie destruction
+Person1: I know we blew up that bank
+Person3: Ah sense negative
+Person1: let's not In know you put-
+Person2: Do you understand it or do you find it confusing?
+Person3: Oh! Le blue fance!
+Person1: Oh my goodness the police!! Get down! It's the police!
+Person2: She glanced out of the window
+Person1: No No! Their going past just keep quiet!
+Person2: MONEY!
+Person3: VIOLENCE CRIMINEL!
+Person1: NO NO! SHUSH! 
+Person2: I EXPECT YOU'LL BE WANTING SOME CHANGE
+Person1: Listen, listen, guys keep your f***ing voices down.
+Person2: MONEYYYYYY!!!!
+Person1: Okay I think they've gone. You want money? Fine here's your money
+Person2: Seventy Seven
+Person1: Seventy Seven pounds. Here you go. Congratulations. And for you-
+Person2: 77!?
+Person1: That's what we agreed-
+Person3: TYRANNY
+Person2: 77!?
+Person1: Guys. Listen it wasn't the biggest bank robbery in the world. We robbed 100 quid. Alright? You get 77
+Person2: Money?
+Person1: I'll get 10 and Person3 gets whatever the rest of the maths is.
+Person3: Uhhhh... 108.
+**End**

__init__.py

@@ -0,0 +1,64 @@
+import os
+import time
+import torch
+from transformers import GPT2LMHeadModel, GPT2Tokenizer, TextDataset, DataCollatorForLanguageModeling, Trainer, TrainingArguments
+
+class GptHumorTrainer:
+
+    def __init__(self, silent=False) -> None:
+        start_time = time.perf_counter()
+        self.tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
+        self.model = GPT2LMHeadModel.from_pretrained(self.local_file_path("SaveState"))
+        self.model.eval()
+        if not silent:
+            print(f"Model Loading Took {time.perf_counter()-start_time} Seconds")
+
+    def local_file_path(self, path):
+        return os.path.join(os.path.dirname(os.path.abspath(__file__)), path)
+
+    def train(self, train_file, epochs=3):
+        device = torch.device("cpu")
+        self.model.to(device)
+
+        # Prepare the dataset
+        train_dataset = TextDataset(
+            tokenizer=self.tokenizer,
+            file_path=train_file,
+            block_size=128,
+        )
+
+        # We use a special data collator for language modeling tasks
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=self.tokenizer,
+            mlm=False,
+        )
+
+        for epoch in range(epochs):
+            # Define the training arguments for each epoch
+            training_args = TrainingArguments(
+                output_dir=f"./results/epoch_{epoch+1}",  # The output directory for this epoch
+                overwrite_output_dir=True,  # Overwrite the content of the output directory
+                num_train_epochs=3,  # Train for 1 epoch at a time
+                per_device_train_batch_size=3,  # Batch size for training
+                save_steps=-1,  # Save model after each epoch
+                save_total_limit=None,  # No limit on the total amount of checkpoints
+                prediction_loss_only=True,  # Focus on the prediction loss only
+            )
+
+            # Initialize the Trainer
+            trainer = Trainer(
+                model=self.model,
+                args=training_args,
+                data_collator=data_collator,
+                train_dataset=train_dataset,
+            )
+
+            # Train the model for one epoch
+            trainer.train()
+
+            # Save the model after each epoch
+            self.model.save_pretrained(self.local_file_path("SaveState"))
+
+if __name__ == "__main__":
+    humor_trainer = GptHumorTrainer()
+    humor_trainer.train(humor_trainer.local_file_path("TrainData.txt"), epochs=5)  # Replace with the path to your training file