Update app.py

app.py

@@ -2,7 +2,7 @@ import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-MODEL_NAME = "ibm-granite/granite-3.3-2b-instruct"
+MODEL_NAME = "ibm-granite/granite-3.0-2b-base"
 
 # Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
@@ -10,48 +10,37 @@ model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
 )
-model.to("cuda" if torch.cuda.is_available() else "cpu")
-model.eval()
-
-
-def chat(user_input, history):
-    messages = []
 
-    # Convert chat history to Granite chat format
-    for user, assistant in history:
-        messages.append({"role": "user", "content": user})
-        messages.append({"role": "assistant", "content": assistant})
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+model.eval()
 
-    messages.append({"role": "user", "content": user_input})
 
-    inputs = tokenizer.apply_chat_template(
-        messages,
-        add_generation_prompt=True,
-        tokenize=True,
-        return_dict=True,
-        return_tensors="pt",
-    ).to(model.device)
+def generate_text(prompt, max_new_tokens=100, temperature=0.7):
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=100,
+            max_new_tokens=max_new_tokens,
             do_sample=True,
-            temperature=0.7,
+            temperature=temperature,
             top_p=0.9,
         )
 
-    response = tokenizer.decode(
-        outputs[0][inputs["input_ids"].shape[-1]:],
-        skip_special_tokens=True,
-    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    history.append((user_input, response))
-    return history, history
 
+demo = gr.Interface(
+    fn=generate_text,
+    inputs=[
+        gr.Textbox(lines=5, label="Input Prompt"),
+        gr.Slider(10, 300, value=100, step=10, label="Max New Tokens"),
+        gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature"),
+    ],
+    outputs=gr.Textbox(lines=10, label="Generated Output"),
+    title="IBM Granite 3.0 – 2B Base",
+    description="Text generation using IBM Granite 3.0 2B Base model",
+)
 
-gr.ChatInterface(
-    fn=chat,
-    title="IBM Granite 3.3 2B Instruct",
-    description="Chat with IBM Granite using Hugging Face Transformers",
-).launch()
+demo.launch()