iShare's picture
Update main.py
0f7b3d4
raw
history blame
2.57 kB
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from langchain.document_loaders import WebBaseLoader
from langchain.chains.summarize import load_summarize_chain
#https://python.langchain.com/docs/use_cases/summarization
from langchain import HuggingFaceHub
from huggingface_hub import InferenceClient
from bs4 import BeautifulSoup
import requests
import sys
import os
from dotenv import load_dotenv
load_dotenv()
hf_token = os.environ.get('HUGGINGFACEHUB_API_TOKEN')
repo_id=os.environ.get('repo_id')
#port = os.getenv('port')
llm = HuggingFaceHub(repo_id=repo_id, # for StarChat
huggingfacehub_api_token=hf_token,
model_kwargs={"min_length": 512, # for StarChat
"max_new_tokens": 1024, "do_sample": True, # for StarChat
"temperature": 0.01,
"top_k": 50,
"top_p": 0.95, "eos_token_id": 49155})
#chain = load_summarize_chain(llm, chain_type="stuff")
chain = load_summarize_chain(llm, chain_type="refine")
app = FastAPI()
class ChatRequest(BaseModel):
target_url: str
#似乎必须要有这个class定义???
#@app.post('/')
#async def home_api(request: Request):
# data = await request.json()
# target_url = data['target_url']
# print(target_url)
# return {"Message": "FastAPI Home API Deploy Success on HF"}
@app.post('/api/chat')
async def chat(request: Request, chat_request: ChatRequest):
target_url = chat_request.target_url
#@app.post('/api/chat')
#async def chat(request: Request):
# data = await request.json()
# target_url = data['target_url']
#以上四行代码,可能有问题?出在data = await request.json()???理论上似乎也没有问题,因为后面的print(target_url),能够成功打印
print(target_url)
try:
loader = WebBaseLoader(target_url)
print(target_url)
docs = loader.load()
result = chain.run(docs)
print("AI Summarization: " + result)
#return {'response': result} #FastAPI方式下,这个返回形式,有问题?
return JSONResponse({'response': result})
#对比FlaskAPI的:return jsonify({'response': result}),JSONResponse vs jsonify(都需要相应的from ... import ...)
except Exception as e:
err_msg = "Wrong URL or URL not parsable."
print(err_msg)
raise HTTPException(status_code=400, detail=err_msg)