yashxx07's picture
Update main.py
7cfa396 verified
raw
history blame
2.65 kB
import requests
from bs4 import BeautifulSoup
from fastapi import FastAPI
from pydantic import BaseModel
import re
import os
from transformers import AutoModelWithLMHead, AutoTokenizer
app = FastAPI()
tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-summarize-news", use_fast=False)
model = AutoModelWithLMHead.from_pretrained("mrm8488/t5-base-finetuned-summarize-news", use_fast=False)
def summarize(text, max_length=150):
input_ids = tokenizer.encode(text, return_tensors="pt", add_special_tokens=True)
generated_ids = model.generate(input_ids=input_ids, num_beams=2, max_length=max_length, repetition_penalty=2.5, length_penalty=1.0, early_stopping=True)
preds = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g in generated_ids]
return preds[0]
@app.get("/")
async def root():
return {"status": "OK"}
class Item(BaseModel):
url: str
percentage: int
def extract_article_content(url):
try:
# Fetch the HTML content of the article URL
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.post("/summarize-v2")
async def root(item: Item):
try:
article = extract_article_content(item.url)
if len(article) == 0:
return {'summary': ""}
response = requests.post('https://fumes-api.onrender.com/llama3',
json={'prompt': "{ 'User': 'Summarize the following news article: '" + article + "}",
"temperature":0.6,
"topP":0.9,
"maxTokens": 200}, stream=True)
response_content = response.content.decode('utf-8')
response_content = response_content.replace("Here is a summary of the news article:", "")
response_content = response_content.replace("YOU CAN BUY ME COFFE! https://buymeacoffee.com/mygx", "")
#return {clean_response}
return {
"summary":response_content}
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}