Summirizer_agent Update

main.py

@@ -1,12 +1,13 @@
 from fastapi import FastAPI, Form
 from fastapi.responses import HTMLResponse
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 
 app = FastAPI()
 
-# Load Granite 2B model
 MODEL_ID = "ibm-granite/granite-3.3-2b-instruct"
+
+# Load tokenzier and model
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
@@ -14,6 +15,9 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 
+# Use pipeline for easier text generation
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
+
 @app.get("/", response_class=HTMLResponse)
 def index():
     return """
@@ -31,15 +35,21 @@ def index():
 
 @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,
+    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:"
+    )
+    # Generate output using the pipeline
+    outputs = pipe(
+        prompt,
         max_new_tokens=150,
-        do_sample=False,
-        temperature=0.7
+        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 = 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>"
+    output_text = outputs[0]['generated_text']
+    # Extract only the summary after 'Summary:'
+    summary = output_text.split("Summary:")[-1].strip()
+    return f"<h2>Summary</h2><pre>{summary}</pre><a href='/'>Back</a>"