BeveledCube commited on
Commit
6b6d0ef
·
verified ·
1 Parent(s): 8e74647

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +9 -2
main.py CHANGED
@@ -9,12 +9,19 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer
9
  import torch
10
 
11
  app = FastAPI()
12
- name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
13
- # mistralai/Mixtral-8x7B-Instruct-v0.1
14
  # microsoft/DialoGPT-small
15
  # microsoft/DialoGPT-medium
16
  # microsoft/DialoGPT-large
17
 
 
 
 
 
 
 
 
18
  # Load the Hugging Face GPT-2 model and tokenizer
19
  model = GPT2LMHeadModel.from_pretrained(name)
20
  tokenizer = GPT2Tokenizer.from_pretrained(name)
 
9
  import torch
10
 
11
  app = FastAPI()
12
+ name = "PygmalionAI/pygmalion-2.7b"
13
+
14
  # microsoft/DialoGPT-small
15
  # microsoft/DialoGPT-medium
16
  # microsoft/DialoGPT-large
17
 
18
+ # PygmalionAI/pygmalion-350m
19
+ # PygmalionAI/pygmalion-1.3b
20
+ # PygmalionAI/pygmalion-2.7b
21
+ # PygmalionAI/pygmalion-6b
22
+
23
+ # mistralai/Mixtral-8x7B-Instruct-v0.1
24
+
25
  # Load the Hugging Face GPT-2 model and tokenizer
26
  model = GPT2LMHeadModel.from_pretrained(name)
27
  tokenizer = GPT2Tokenizer.from_pretrained(name)