|
|
from fastapi import FastAPI, Form |
|
|
from fastapi.responses import HTMLResponse |
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline |
|
|
import torch |
|
|
|
|
|
app = FastAPI() |
|
|
|
|
|
MODEL_ID = "ibm-granite/granite-4.0-tiny-preview" |
|
|
|
|
|
|
|
|
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" |
|
|
) |
|
|
|
|
|
|
|
|
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer) |
|
|
|
|
|
@app.get("/", response_class=HTMLResponse) |
|
|
def index(): |
|
|
return """ |
|
|
<html> |
|
|
<head><title>Granite Tiny Summarizer</title></head> |
|
|
<body> |
|
|
<h1>Granite 4.0 Tiny 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 = ( |
|
|
"Below is a passage of text. Please provide a concise summary in 2-4 sentences.\n\n" |
|
|
f"Text:\n{text.strip()}\n\nSummary:" |
|
|
) |
|
|
outputs = pipe( |
|
|
prompt, |
|
|
max_new_tokens=150, |
|
|
do_sample=True, |
|
|
temperature=0.7, |
|
|
top_p=0.95, |
|
|
eos_token_id=tokenizer.eos_token_id, |
|
|
pad_token_id=tokenizer.eos_token_id |
|
|
) |
|
|
output_text = outputs[0]['generated_text'] |
|
|
summary = output_text.split("Summary:")[-1].strip() |
|
|
return f"<h2>Summary</h2><pre>{summary}<br></pre><a href='/'>Back</a>" |
