syedmoinms commited on
Commit
055aed6
·
verified ·
1 Parent(s): fc72167

Update server.py

Browse files
Files changed (1) hide show
  1. server.py +3 -3
server.py CHANGED
@@ -4,9 +4,9 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
4
  app = FastAPI()
5
 
6
  # Load Fine-Tuned Model
7
- model_name = "/app/models/fine-tuned"
8
- tokenizer = AutoTokenizer.from_pretrained("NousResearch/Nous-Hermes-2-Mistral-7B")
9
- model = AutoModelForCausalLM.from_pretrained("NousResearch/Nous-Hermes-2-Mistral-7B")
10
 
11
  @app.get("/chat")
12
  def chat(msg: str):
 
4
  app = FastAPI()
5
 
6
  # Load Fine-Tuned Model
7
+ model_name = "PygmalionAI/pygmalion-7b"
8
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
9
+ model = AutoModelForCausalLM.from_pretrained(model_name)
10
 
11
  @app.get("/chat")
12
  def chat(msg: str):