test

app.py

@@ -23,7 +23,7 @@ from dotenv import load_dotenv
 
 zero = torch.Tensor([0]).cuda()
 
-
+@spaces.GPU
 def read_pdf(file_path):
     output = ''
     doc = pymupdf.open(file_path)
@@ -49,34 +49,34 @@ def read_pdf(file_path):
 
 # Function to query Hugging Face endpoint
 #@spaces.GPU
-def query_huggingface(text):
-    print(zero.device)
-    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])
+# def query_huggingface(text):
+#     print(zero.device)
+#     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):
     pdf_output = read_pdf(file.name)
-    huggingface_output = query_huggingface(query)
-    return pdf_output, huggingface_output
+    #huggingface_output = query_huggingface(query)
+    return pdf_output#, huggingface_output
 
 # Create Gradio App
 interface = gr.Interface(
     fn=process_file,
     inputs=[
         gr.File(label="Upload a PDF file"),
-        gr.Textbox(label="Enter your query for Hugging Face"),
+        #gr.Textbox(label="Enter your query for Hugging Face"),
     ],
     outputs=[
         gr.Textbox(label="PDF Content"),
-        gr.Textbox(label="Hugging Face Output"),
+       #gr.Textbox(label="Hugging Face Output"),
     ],
     title="PDF Processor with Hugging Face Query"
 )