yashxx07's picture
Update main.py
09c5387 verified
raw
history blame
2.65 kB
import requests
from bs4 import BeautifulSoup
from fastapi import FastAPI#, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import re
import replicate
class Item(BaseModel):
url: str
max_tokens: int
app = FastAPI()
def extract_article_content(url):
try:
response = requests.get(url)
soup = BeautifulSoup(response.text, 'html.parser')
results = soup.find_all(['h1', 'p'])
text = [result.text for result in results]
ARTICLE = ' '.join(text)
return ARTICLE
except Exception as e:
return ""
@app.get("/")
async def root():
return {"status": "OK"}
@app.post("/summarize-v1")
async def root(item: Item):
try:
article = extract_article_content(item.url)
if len(article) == 0:
return {'summary': ""}
event_list = []
for event in replicate.stream("snowflake/snowflake-arctic-instruct", input= {
"prompt": "summarize this following news article:" + article,
"temperature": 0.2}):
# Convert the event to a string and append it to the list
event_list.append(str(event))
# After the event stream ends, process the collected events
output_variable = "".join(event_list)
return output_variable
except requests.RequestException as e:
return {"error": str(e), "status_code": 500}
@app.post("/summarize-v2")
async def root(item: Item):
try:
article = extract_article_content(item.url)
if len(article) == 0:
return {'summary': ""}
def event_generator():
for event in replicate.stream("snowflake/snowflake-arctic-instruct", input={
"prompt": "summarize this following news article:" + article,
"temperature": 0.2,
"max_new_tokens" : item.max_tokens
}):
# Yield the event as a string
yield str(event)
print(str(event), end="")
# Use StreamingResponse to stream the events
return StreamingResponse(event_generator(), media_type='text/event-stream')
except requests.RequestException as e:
return {"error": str(e), "status_code": 500}
@app.post("/extract-content")
async def root(item: Item):
try:
article = extract_article_content(item.url)
if len(article) == 0:
return {'ERROR': "AHHHHHHHHH"}
return {"content":article}
except requests.RequestException as e:
return {"error": str(e), "status_code": 500}