creating a custom wrapper to move inputs to gpu

app.py

@@ -3,35 +3,37 @@ import torch
 from transformers import (
     AutoTokenizer,
     AutoModelForCausalLM,
-    pipeline,
     BitsAndBytesConfig,
+    pipeline
 )
 from langchain_experimental.agents import create_pandas_dataframe_agent
 from langchain_community.llms import HuggingFacePipeline
 import gradio as gr
-import spaces  # required for ZeroGPU
+import spaces
 
-# --- Constants ---
-LLM_MODEL_ID    = "HuggingFaceH4/zephyr-7b-beta"
-DATA_FILE       = "IPL.csv"
-MAX_NEW_TOKENS  = 256
-GPU_DURATION    = 120  # seconds for @spaces.GPU
 
-# --- 1) Load & prepare DataFrame once ---
-def load_data():
+# --- Config ---
+LLM_MODEL_ID = "HuggingFaceH4/zephyr-7b-beta"
+DATA_FILE = "IPL.csv"
+
+
+# --- Load IPL Data ---
+def load_df():
     df = pd.read_csv(DATA_FILE, low_memory=False)
     df.columns = df.columns.str.replace(" ", "_").str.lower()
     if "date" in df.columns:
         df["date"] = pd.to_datetime(df["date"], errors="coerce")
-    if all(c in df.columns for c in ("runs_batter", "runs_extras")):
+    if {"runs_batter", "runs_extras"}.issubset(df.columns):
         df["runs_batter"] = pd.to_numeric(df["runs_batter"], errors="coerce").fillna(0)
-        df["runs_extras"]  = pd.to_numeric(df["runs_extras"], errors="coerce").fillna(0)
+        df["runs_extras"] = pd.to_numeric(df["runs_extras"], errors="coerce").fillna(0)
         df["total_runs_this_ball"] = df["runs_batter"] + df["runs_extras"]
     return df
 
-_df = load_data()
 
-# --- 2) Instantiate tokenizer, model, pipeline, and agent globally ---
+_df = load_df()
+
+
+# --- Load Quantized Model ---
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
@@ -40,80 +42,99 @@ bnb_config = BitsAndBytesConfig(
 )
 
 tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_ID, trust_remote_code=True)
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
+tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token
 
 model = AutoModelForCausalLM.from_pretrained(
     LLM_MODEL_ID,
-    quantization_config=bnb_config,
     torch_dtype=torch.float16,
+    device_map="auto",
+    quantization_config=bnb_config,
     trust_remote_code=True,
 )
-# model.to("cuda")
-
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    # device=0,                  # <— ensure GPU inference
-    max_new_tokens=MAX_NEW_TOKENS,
-    do_sample=True,
-    temperature=0.1,
-    top_p=0.9,
-    eos_token_id=tokenizer.eos_token_id,
-    pad_token_id=tokenizer.pad_token_id,
-)
 
-hf_llm = HuggingFacePipeline(pipeline=pipe)
-# (NO hf_llm.to("cuda"); the pipeline already handles device)
 
+# --- LLM Wrapper for LangChain ---
+class MyLLMWrapper:
+    def __init__(self):
+        self.tokenizer = tokenizer
+        self.model = model
+
+    def invoke(self, input_str):
+        return self.__call__(input_str)
+
+    def __call__(self, input_str):
+        inputs = self.tokenizer(input_str, return_tensors="pt").to(self.model.device)
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=256,
+                do_sample=True,
+                temperature=0.1,
+                top_p=0.9,
+                eos_token_id=self.tokenizer.eos_token_id,
+                pad_token_id=self.tokenizer.pad_token_id
+            )
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+
+
+llm = MyLLMWrapper()
+
+
+# --- System Prompt for the Agent ---
 system_message = """
-You are an expert cricket analyst. You have access to a pandas DataFrame named `df`
-containing ball-by-ball IPL match data. Use Python (pandas) to answer queries about IPL stats
-as efficiently as possible. Do not import extra libraries.
+You are an expert IPL cricket analyst. You have access to a pandas DataFrame named `df` that contains ball-by-ball IPL match data.
+Answer all questions using pandas logic, match stats, and accurate calculations.
 """
+
+# --- LangChain Agent ---
 agent = create_pandas_dataframe_agent(
-    hf_llm,
+    llm,
     _df,
     verbose=False,
-    max_iterations=5,
     handle_parsing_errors=True,
     agent_executor_kwargs={"system_message": system_message},
-    agent_type="openai-tools",
+    agent_type="openai-tools",  # Most compatible with Hugging Face models
     allow_dangerous_code=True,
 )
 
-# --- 3) Define inference function (GPU-enabled) ---
-@spaces.GPU(duration=GPU_DURATION)
-def run_inference(question: str) -> str:
-    torch.cuda.empty_cache()  # free up cached memory
-    result = agent.invoke({"input": question})
-    return result.get("output", "No output returned.")
 
-# --- 4) Build Gradio app ---
-def bot_response(history):
-    query = history[-1][0]
+# --- Inference Function ---
+@spaces.GPU(duration=120)
+def predict_answer(question):
+    torch.cuda.empty_cache()
     try:
-        answer = run_inference(query)
+        res = agent.invoke({"input": question})
+        return res.get("output", "No response generated.")
     except Exception as e:
-        answer = f"Error during inference: {e}"
-    history[-1][1] = answer
-    return history
+        return f"❌ Error during inference: {e}"
+
 
+# --- Gradio UI ---
 with gr.Blocks() as demo:
-    gr.Markdown("# IPL Cricket Data Agent")
-    gr.Markdown("Ask me anything about the IPL dataset (e.g., top run-scorers, match outcomes, averages).")
-    chatbot = gr.Chatbot()
-    user_input = gr.Textbox(placeholder="Type your question here...")
-    clear_btn  = gr.Button("Clear")
-
-    user_input.submit(
-        lambda msg, chat: (None, chat + [[msg, None]]),
-        [user_input, chatbot],
-        [user_input, chatbot],
-        queue=True
-    ).then(bot_response, chatbot, chatbot)
-
-    clear_btn.click(lambda: [], None, chatbot)
+    gr.Markdown("# 🏏 IPL Cricket Analyst")
+    gr.Markdown(
+        "Ask questions about IPL stats from the dataset. Examples:<br>"
+        "`Top 5 batsmen by total runs`<br>"
+        "`Who scored the most in 2023?`<br>"
+        "`Average runs per over in 2022?`"
+    )
+
+    chatbot = gr.Chatbot(label="Cricket Analyst")
+    msg = gr.Textbox(label="Ask your question here...")
+    clear = gr.Button("Clear")
+
+    def user_input(m, hist):
+        return "", hist + [[m, None]]
+
+    def bot_reply(hist):
+        q = hist[-1][0]
+        a = predict_answer(q)
+        hist[-1][1] = a
+        return hist
+
+    msg.submit(user_input, [msg, chatbot], [msg, chatbot], queue=True).then(
+        bot_reply, chatbot, chatbot
+    )
+    clear.click(lambda: [], None, chatbot)
 
 demo.queue(max_size=20).launch(debug=True)