Update app.py

app.py

@@ -1,89 +1,27 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
 import gradio as gr
-from fastapi import FastAPI, Request
-
-MODEL_ID = "ibm-granite/granite-3.1-2b-instruct"
-
-# Load model + tokenizer
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto"
-)
-
-DEFAULT_PARAMS = dict(
-    max_new_tokens=512,
-    temperature=0.2,
-    top_p=0.95,
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+
+model_id = "ibm-granite/granite-3.1-2b-instruct"  # Hugging Face model ID
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
+
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_length=512,
+    temperature=0.7
 )
 
-def format_instruct_prompt(system_msg, user_msg):
-    # Works with Granite chat/instruct style
-    # Adjust if your prompt format differs
-    return f"<|system|>\n{system_msg}\n<|user|>\n{user_msg}\n<|assistant|>\n"
-
-def generate_once(system, user, params=None):
-    if params is None:
-        params = {}
-    merged = {**DEFAULT_PARAMS, **params}
-    prompt = format_instruct_prompt(system, user)
-
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    output_ids = model.generate(
-        **inputs,
-        max_new_tokens=merged["max_new_tokens"],
-        do_sample=merged["temperature"] > 0,
-        temperature=merged["temperature"],
-        top_p=merged["top_p"],
-        pad_token_id=tokenizer.eos_token_id,
-        eos_token_id=tokenizer.eos_token_id,
-    )
+def generate_answer(prompt):
+    result = pipe(prompt)[0]["generated_text"]
+    return result
 
-    # Only decode the newly generated tokens
-    gen_ids = output_ids[0][inputs["input_ids"].shape[-1]:]
-    text = tokenizer.decode(gen_ids, skip_special_tokens=True)
-    return text.strip()
-
-# ---------- Gradio UI (manual testing) ----------
-with gr.Blocks() as demo:
-    gr.Markdown("# Granite RAG API (UI)")
-    sys_in = gr.Textbox(label="System", value="You are Granite, a helpful and concise assistant.")
-    usr_in = gr.Textbox(label="User", placeholder="Ask something...")
-    max_new = gr.Slider(64, 1024, value=512, step=16, label="max_new_tokens")
-    temperature = gr.Slider(0.0, 1.5, value=0.2, step=0.05, label="temperature")
-    top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.01, label="top_p")
-    out = gr.Textbox(label="Output")
-
-    def _ui(system, user, max_new_tokens, temperature, top_p):
-        return generate_once(system, user, {
-            "max_new_tokens": int(max_new_tokens),
-            "temperature": float(temperature),
-            "top_p": float(top_p),
-        })
-
-    gr.Button("Generate").click(_ui, [sys_in, usr_in, max_new, temperature, top_p], out)
-
-# ---------- FastAPI JSON endpoint ----------
-api = FastAPI()
-
-@api.post("/generate")
-async def generate(req: Request):
-    """
-    POST JSON:
-    {
-      "prompt": "question with context...",
-      "system": "system prompt (optional)",
-      "params": { "max_new_tokens": 300, "temperature": 0.2, "top_p": 0.9 }
-    }
-    """
-    body = await req.json()
-    prompt = body["prompt"]
-    system = body.get("system", "You are Granite, a helpful and concise assistant.")
-    params = body.get("params", {})
-    text = generate_once(system, prompt, params)
-    return {"text": text}
+demo = gr.Interface(
+    fn=generate_answer,
+    inputs="text",
+    outputs="text",
+    title="Granite 3.1 2B Instruct - RAG Answering"
+)
 
-# Mount Gradio on "/"
-app = gr.mount_gradio_app(api, demo, path="/")
+demo.launch()