test

app.py

@@ -7,7 +7,7 @@ import torch
 import gradio as gr
 from prepare import prepare
 
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from langchain_community.llms import HuggingFacePipeline
 from langchain.prompts import PromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate
 from langchain_core.output_parsers import StrOutputParser
@@ -51,17 +51,15 @@ def read_pdf(file_path):
 @spaces.GPU
 def query_huggingface(text):
     print(zero.device)
-    pipe = pipeline(
-    "text-generation",
-    model="google/gemma-2-9b-it",
-    model_kwargs={"torch_dtype": torch.bfloat16},
-    device="cuda",  # replace with "mps" to run on a Mac device
-)
-    messages = [
-    {"role": "user", "content": text},
-]
-    outputs = pipe(messages, max_new_tokens=256)
-    return outputs[0]["generated_text"][-1]["content"].strip()
+    tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b-it")
+    model = AutoModelForCausalLM.from_pretrained(
+    "google/gemma-2-9b-it",
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+        )
+    input_ids = tokenizer(text, return_tensors="pt").to("cuda")
+    outputs = model.generate(**input_ids, max_new_tokens=32)
+    return tokenizer.decode(outputs[0])
 
 # Gradio Interface for PDF Processing
 def process_file(file, query):