test

app.py

@@ -7,7 +7,7 @@ import torch
 import gradio as gr
 from prepare import prepare
 
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
 from langchain_community.llms import HuggingFacePipeline
 from langchain.prompts import PromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate
 from langchain_core.output_parsers import StrOutputParser
@@ -51,16 +51,17 @@ def read_pdf(file_path):
 @spaces.GPU
 def query_huggingface(text):
     print(zero.device)
-    load_dotenv()
-    api_token = os.getenv("API_TOKEN")
-    repo_id = "google/gemma-2-9b-it"
-    task = "text-generation"
-    chat_model = HuggingFaceEndpoint(
-        huggingfacehub_api_token=api_token,
-        repo_id=repo_id,
-        task=task
-    )
-    return chat_model.invoke(text)
+    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()
 
 # Gradio Interface for PDF Processing
 def process_file(file, query):