|
from fastapi import FastAPI, HTTPException |
|
from typing import Any, Dict |
|
from pydantic import BaseModel |
|
import os |
|
from os import getenv |
|
from huggingface_hub import InferenceClient |
|
import random |
|
from json_repair import repair_json |
|
import nltk |
|
|
|
app = FastAPI() |
|
|
|
nltk.download('punkt') |
|
|
|
tokenizer = nltk.data.load('tokenizers/punkt/english.pickle') |
|
|
|
HF_TOKEN = getenv("HF_TOKEN") |
|
|
|
class InputData(BaseModel): |
|
model: str |
|
system_prompt_template: str |
|
prompt_template: str |
|
end_token: str |
|
system_prompt: str |
|
user_input: str |
|
json_prompt: str |
|
history: str = "" |
|
|
|
@app.post("/generate-response/") |
|
async def generate_response(data: InputData) -> Dict[str, Any]: |
|
client = InferenceClient(model=data.model, token=HF_TOKEN) |
|
|
|
sentences = tokenizer.tokenize(data.user_input) |
|
data_dict = {'###New response###': [], '###Sentence count###': 0} |
|
for i, sentence in enumerate(sentences): |
|
data_dict["###New response###"].append(sentence) |
|
data_dict["###Sentence count###"] = i + 1 |
|
|
|
data.history += data.prompt_template.replace("{Prompt}", str(data_dict)) |
|
|
|
inputs = ( |
|
data.system_prompt_template.replace("{SystemPrompt}", data.system_prompt) + |
|
data.system_prompt_template.replace("{SystemPrompt}", data.json_prompt) + |
|
data.history |
|
) |
|
|
|
seed = random.randint(0, 2**32 - 1) |
|
|
|
try: |
|
response = client.text_generation( |
|
inputs, |
|
temperature=1.0, |
|
max_new_tokens=1000, |
|
seed=seed |
|
) |
|
|
|
strict_response = str(response) |
|
|
|
repaired_response = repair_json(strict_response, return_objects=True) |
|
|
|
if isinstance(repaired_response, str): |
|
raise HTTPException(status_code=500, detail="Invalid response from model") |
|
else: |
|
cleaned_response = {} |
|
for key, value in repaired_response.items(): |
|
cleaned_key = key.replace("###", "") |
|
cleaned_response[cleaned_key] = value |
|
|
|
for i, text in enumerate(cleaned_response["New response"]): |
|
if i <= 2: |
|
sentences = tokenizer.tokenize(text) |
|
if sentences: |
|
cleaned_response["New response"][i] = sentences[0] |
|
else: |
|
del cleaned_response["New response"][i] |
|
|
|
if cleaned_response.get("Sentence count"): |
|
if cleaned_response["Sentence count"] > 3: |
|
cleaned_response["Sentence count"] = 3 |
|
else: |
|
cleaned_response["Sentence count"] = len(cleaned_response["New response"]) |
|
|
|
data.history += str(cleaned_response) |
|
|
|
return { |
|
"response": cleaned_response, |
|
"history": data.history + data.end_token |
|
} |
|
|
|
except Exception as e: |
|
print(f"Model {data.model} failed with error: {e}") |
|
raise HTTPException(status_code=500, detail=f"Model {data.model} failed to generate response") |
|
|
|
@app.post("/get-medieval-name/") |
|
async def get_medieval_name() -> Dict[str, str]: |
|
try: |
|
file_path = "medieval_names.txt" |
|
if not os.path.exists(file_path): |
|
raise HTTPException(status_code=404, detail="File not found") |
|
|
|
with open(file_path, "r") as file: |
|
names = file.read().splitlines() |
|
|
|
if not names: |
|
raise HTTPException(status_code=404, detail="No names found in the file") |
|
|
|
random_name = random.choice(names) |
|
|
|
return {"name": random_name} |
|
|
|
except Exception as e: |
|
print(f"Error: {e}") |
|
raise HTTPException(status_code=500, detail="An error occurred while processing the request") |
