test

app.py

@@ -2,7 +2,7 @@ from PIL import Image
 import pytesseract
 import os
 import pymupdf
-#import spaces
+import spaces
 import torch
 import gradio as gr
 from prepare import prepare
@@ -21,7 +21,7 @@ from langchain_core.messages import AIMessage, HumanMessage
 from langchain_community.llms import HuggingFaceEndpoint
 from dotenv import load_dotenv
 from huggingface_hub import InferenceClient
-
+import huggingface_hub
 #zero = torch.Tensor([0]).cuda()
 load_dotenv()
 api_token = os.getenv("HF_TOKEN")
@@ -53,43 +53,57 @@ def read_pdf(file_path):
 
 # Function to query Hugging Face endpoint
 #@spaces.GPU
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+# def respond(
+#     message,
+#     history: list[tuple[str, str]],
+#     system_message,
+#     max_tokens,
+#     temperature,
+#     top_p,
+# ):
+#     messages = [{"role": "system", "content": system_message}]
+
+#     for val in history:
+#         if val[0]:
+#             messages.append({"role": "user", "content": val[0]})
+#         if val[1]:
+#             messages.append({"role": "assistant", "content": val[1]})
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+#     messages.append({"role": "user", "content": message})
 
-    messages.append({"role": "user", "content": message})
+#     response = ""
 
-    response = ""
+#     for message in client.chat_completion(
+#         messages,
+#         max_tokens=max_tokens,
+#         stream=True,
+#         temperature=temperature,
+#         top_p=top_p,
+#     ):
+#         token = message.choices[0].delta.content
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+#         response += token
+#     return response
 
-        response += token
-    return response
+@spaces.GPU
+def LLM_Inference(text):
+    huggingface_hub.login(token=api_token)
+    device = torch.device('cuda')
+    tokenizer = AutoTokenizer.from_pretrained('google/gemma-2-2b-it')
+    model = AutoModelForCausalLM.from_pretrained('google/gemma-2-2b-it').to(device)
+    inputs = tokenizer(text, return_tensors='pt').to(device)
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs, max_new_tokens=128, pad_token_id = tokenizer.eos_token_id
+        )
+    return tokenizer.decode(outputs[0])
 
 # Gradio Interface for PDF Processing
 def process_file(file, query):
     pdf_output = read_pdf(file.name)
-    huggingface_output = respond(query, history=[], system_message="You are a friendly Chatbot.",
-                                max_tokens=1024, temperature=0.0, top_p=0.99 )
+    #huggingface_output = respond(query, history=[], system_message="You are a friendly sChatbot.",
+                                #max_tokens=1024, temperature=0.0, top_p=0.99 )
+    huggingface_output = LLM_Inference(query)
     return pdf_output, huggingface_output
 
 # Create Gradio App