Testing granite-4.0-tiny-preview
Browse files
main.py
CHANGED
|
@@ -5,7 +5,7 @@ import torch
|
|
| 5 |
|
| 6 |
app = FastAPI()
|
| 7 |
|
| 8 |
-
MODEL_ID = "ibm-granite/granite-
|
| 9 |
|
| 10 |
# Load tokenizer and model
|
| 11 |
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
|
@@ -15,16 +15,16 @@ model = AutoModelForCausalLM.from_pretrained(
|
|
| 15 |
device_map="auto"
|
| 16 |
)
|
| 17 |
|
| 18 |
-
# Use pipeline for easier text generation (device
|
| 19 |
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
|
| 20 |
|
| 21 |
@app.get("/", response_class=HTMLResponse)
|
| 22 |
def index():
|
| 23 |
return """
|
| 24 |
<html>
|
| 25 |
-
<head><title>Granite
|
| 26 |
<body>
|
| 27 |
-
<h1>Granite
|
| 28 |
<form action="/summarize" method="post">
|
| 29 |
<textarea name="text" rows="10" cols="80" placeholder="Paste text to summarize"></textarea><br>
|
| 30 |
<button type="submit">Summarize</button>
|
|
@@ -39,7 +39,6 @@ def summarize(text: str = Form(...)):
|
|
| 39 |
"Below is a passage of text. Please provide a concise summary in 2-4 sentences.\n\n"
|
| 40 |
f"Text:\n{text.strip()}\n\nSummary:"
|
| 41 |
)
|
| 42 |
-
# Generate output using the pipeline
|
| 43 |
outputs = pipe(
|
| 44 |
prompt,
|
| 45 |
max_new_tokens=150,
|
|
@@ -50,6 +49,5 @@ def summarize(text: str = Form(...)):
|
|
| 50 |
pad_token_id=tokenizer.eos_token_id
|
| 51 |
)
|
| 52 |
output_text = outputs[0]['generated_text']
|
| 53 |
-
# Extract only the summary after 'Summary:'
|
| 54 |
summary = output_text.split("Summary:")[-1].strip()
|
| 55 |
-
return f"<h2>Summary</h2><pre>{summary}
|
|
|
|
| 5 |
|
| 6 |
app = FastAPI()
|
| 7 |
|
| 8 |
+
MODEL_ID = "ibm-granite/granite-4.0-tiny-preview"
|
| 9 |
|
| 10 |
# Load tokenizer and model
|
| 11 |
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
|
|
|
| 15 |
device_map="auto"
|
| 16 |
)
|
| 17 |
|
| 18 |
+
# Use pipeline for easier text generation (no device arg when using device_map="auto")
|
| 19 |
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
|
| 20 |
|
| 21 |
@app.get("/", response_class=HTMLResponse)
|
| 22 |
def index():
|
| 23 |
return """
|
| 24 |
<html>
|
| 25 |
+
<head><title>Granite Tiny Summarizer</title></head>
|
| 26 |
<body>
|
| 27 |
+
<h1>Granite 4.0 Tiny Summarization Demo</h1>
|
| 28 |
<form action="/summarize" method="post">
|
| 29 |
<textarea name="text" rows="10" cols="80" placeholder="Paste text to summarize"></textarea><br>
|
| 30 |
<button type="submit">Summarize</button>
|
|
|
|
| 39 |
"Below is a passage of text. Please provide a concise summary in 2-4 sentences.\n\n"
|
| 40 |
f"Text:\n{text.strip()}\n\nSummary:"
|
| 41 |
)
|
|
|
|
| 42 |
outputs = pipe(
|
| 43 |
prompt,
|
| 44 |
max_new_tokens=150,
|
|
|
|
| 49 |
pad_token_id=tokenizer.eos_token_id
|
| 50 |
)
|
| 51 |
output_text = outputs[0]['generated_text']
|
|
|
|
| 52 |
summary = output_text.split("Summary:")[-1].strip()
|
| 53 |
+
return f"<h2>Summary</h2><pre>{summary}<br></pre><a href='/'>Back</a>"
|