diff --git a/main.py b/main.py deleted file mode 100644 index 067f0e3..0000000 --- a/main.py +++ /dev/null @@ -1,53 +0,0 @@ -from fastapi import FastAPI, Request -from fastapi.middleware.cors import CORSMiddleware - -api = FastAPI() - -api.add_middleware( - CORSMiddleware, - allow_origins=['*'], - allow_credentials=True, - allow_methods=['*'], - allow_headers=['*'] - ) - -from transformers.models.auto.tokenization_auto import AutoTokenizer -from transformers.models.auto.modeling_auto import AutoModelForCausalLM - -tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large") -model = AutoModelForCausalLM.from_pretrained("../southpark/output-medium") - -def cartman_respond(packet): - input_ids = tokenizer(str(packet.get('message')) + tokenizer.eos_token, return_tensors="pt").input_ids - outputs = model.generate( - input_ids, - pad_token_id=tokenizer.eos_token_id, - max_new_tokens = int(packet.get('max_new_tokens')), - num_beams = int(packet.get('num_beams')), - num_beam_groups = int(packet.get('num_beam_groups')), - no_repeat_ngram_size = int(packet.get('no_repeat_ngram_size')), - length_penalty = float(packet.get('length_penalty')), - diversity_penalty = float(packet.get('diversity_penalty')), - repetition_penalty = float(packet.get('repetition_penalty')), - early_stopping = bool(packet.get('early_stopping')), - - # do_sample = True, - # top_k = 100, - # top_p = 0.7, - # temperature = 0.8, - ) - return tokenizer.decode(outputs[:, input_ids.shape[-1]:][0], skip_special_tokens=True) - -@api.post('/chat/') -async def getInformation(data : Request): - packet = await data.json() - print(packet) - message = str(packet.get('message')) - print(message) - response = cartman_respond(packet) - print(response) - - return { - "Cartman" : response - } - diff --git a/run b/run deleted file mode 100755 index 6ce80d6..0000000 --- a/run +++ /dev/null @@ -1,3 +0,0 @@ -#!/bin/zsh - -uvicorn main:api --host 10.0.1.1 --reload