File size: 1,489 Bytes
271461a fa4959b 271461a fa4959b 271461a fa4959b 271461a fa4959b 271461a |
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 |
from fastapi import FastAPI, Form
from fastapi.responses import HTMLResponse
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
app = FastAPI()
# Load Granite 2B model
MODEL_ID = "ibm-granite/granite-3.3-2b-instruct"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
torch_dtype=torch.float16 if torch.cuda.is_available() else "auto",
device_map="auto"
)
@app.get("/", response_class=HTMLResponse)
def index():
return """
<html>
<head><title>Granite 2B Summarizer</title></head>
<body>
<h1>Granite 2B Summarization Demo</h1>
<form action="/summarize" method="post">
<textarea name="text" rows="10" cols="80" placeholder="Paste text to summarize"></textarea><br>
<button type="submit">Summarize</button>
</form>
</body>
</html>
"""
@app.post("/summarize", response_class=HTMLResponse)
def summarize(text: str = Form(...)):
prompt = f"Summarize the following text:\n{text.strip()}\nSummary:"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
summary_ids = model.generate(
**inputs,
max_new_tokens=150,
do_sample=False,
temperature=0.7
)
output = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
# Remove the prompt from the output
summary = output.replace(prompt, "").strip()
return f"<h2>Summary</h2><pre>{summary}</pre><a href='/'>Back</a>"
|