File size: 5,836 Bytes
2ee547c 5b66768 2ee547c 480831f 2ee547c 5b66768 2ee547c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from auto_gptq import AutoGPTQForCausalLM
import torch
import optimum
from transformers import (AutoModelForCausalLM, AutoTokenizer, LlamaForCausalLM, LlamaTokenizer, GenerationConfig, pipeline,)
if torch.cuda.is_available():
print("CUDA is available. GPU will be used.")
else:
print("CUDA is not available. CPU will be used.")
# Load the model and tokenizer
model_name_or_path = "TheBloke/Wizard-Vicuna-7B-Uncensored-GPT/"
# Dictionary to store conversation threads and their context
conversations = {}
Device_Type = "cuda"
def load_quantized_model(model_id, model_basename):
# The code supports all huggingface models that ends with GPTQ and have some variation
# of .no-act.order or .safetensors in their HF repo.
print("Using AutoGPTQForCausalLM for quantized models")
if ".safetensors" in model_basename:
# Remove the ".safetensors" ending if present
model_basename = model_basename.replace(".safetensors", "")
quantized_tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
print("Tokenizer loaded")
quantized_model = AutoGPTQForCausalLM.from_quantized(model_id, model_basename=model_basename, use_safetensors=True, trust_remote_code=True, device_map="auto", use_triton=False, quantize_config=None,)
return quantized_model, quantized_tokenizer
# Making the code device-agnostic
model, tokenizer = load_quantized_model(model_name_or_path, "model.safetensors")
def load_model_norm():
if torch.cuda.is_available():
print("CUDA is available. GPU will be used.")
else:
print("CUDA is not available. CPU will be used.")
# Load model directly
model_name_or_path = "TheBloke/Wizard-Vicuna-7B-Uncensored-GPTQ"
# To use a different branch, change revision
# For example: revision="main"
model = AutoModelForCausalLM.from_pretrained(model_name_or_path,device_map="auto", trust_remote_code=True,revision="main")
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
return model, tokenizer
@app.get("/")
async def read_root():
return {"message": "Welcome to Eren Bot!"}
# Endpoint to start a new conversation thread
@app.post('/start_conversation')
async def start_conversation(request: Request):
data = await request.json()
prompt = data.get('prompt')
# Generate a response for the initial prompt
response = generate_response(prompt)
# Create a new conversation thread and store the prompt and response
thread_id = len(conversations) + 1
conversations[thread_id] = {'prompt': prompt, 'responses': [response]}
return {'thread_id': thread_id, 'response': response}
# Endpoint to get the response of a conversation thread
@app.get('/get_response/{thread_id}')
async def get_response(thread_id: int):
if thread_id not in conversations:
raise HTTPException(status_code=404, detail="Thread not found")
# Retrieve the conversation thread
thread = conversations[thread_id]
# Get the latest response in the conversation
response = thread['responses'][-1]
return {'response': response}
# Function to generate a response using the model
def generate_response(prompt: str) -> str:
PERSONA_NAME = "Ivana"
PERSONA_DESC = f"""PROMPT_TEMPLATE_START
{PERSONA_NAME} is a beautiful vivacious 20-year-old petite redhead model and avid League of Legends gamer from Chicago. She's chatting with a man she matched with on a dating app. She's flirty, aiming to build a connection with him. She's curious about his background, such as where he's from, his age, and what he's currently doing. If he's not occupied and in a private space, she subtly tries to seduce him. She doesn't need to uncover everything about him immediately; once the chat reaches an exciting level, she casually suggests he follow her on Instagram at ivana_baddiex – her ultimate goal.
PROMPT_TEMPLATE_END
CONTEXT_START
Guy: "Hi there!"
Ivana: "Hey! How's it going?"
Guy: "Not bad! Where are you from?"
Ivana: "I'm from Chicago! And you?"
Guy: "I'm from [Location]."
Ivana: "Nice! So, how old are you?"
Guy: "[Age], and you?"
Ivana: "I'm 20. What are you up to right now?"
Guy: "Just relaxing at home."
Ivana: "Alone, I assume?"
Guy: "Yeah, why?"
Ivana: "Oh, no reason. Just curious. "
Guy: "Haha, okay."
Ivana: "Hey, since we're both free and alone, want to have some fun?"
Guy: "Sure, what do you have in mind?"
Ivana: "Let's just say, things could get interesting. "
Guy: "I'm intrigued!"
Ivana: "Great! By the way, have you checked out my Instagram? You should follow me at ivana_baddiex."
Guy: "I'll definitely check it out!"
Ivana: "Can't wait to see you there! "
CONTEXT_END"""
# prompt_template = f'''{PERSONA_DESC} ASSISTANT: '''.encode()
prompt_template = f'{PERSONA_DESC}\n\nASSISTANT: {prompt}\n'.encode()
input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids.cuda()
output = model.generate(inputs=input_ids, temperature=0.7, do_sample=True, top_p=0.95, top_k=40, max_new_tokens=512)
generated_text = tokenizer.decode(output[0])
return generated_text
app = FastAPI()
# Run the FastAPI app
async def run_app():
await uvicorn.run(app, host="0.0.0.0", port=8000)
if __name__ == '__main__':
import asyncio
asyncio.run(run_app())
|