bugfix

app.py

@@ -1,4 +1,5 @@
 import os
+# import torch
 import spaces
 import gradio as gr
 from huggingface_hub import login
@@ -23,7 +24,7 @@ def get_llm(model_id):
 @spaces.GPU
 def retriever_qa(file, query):
     model_id = 'mistralai/Mistral-7B-Instruct-v0.2'
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
     llm = get_llm(model_id)
     # retriever_obj = retriever(file)
     # qa = RetrievalQA.from_chain_type(llm=llm, 
@@ -35,15 +36,18 @@ def retriever_qa(file, query):
         first_line = f.readline()
 
     messages = [
-        {"role": "user", "content": first_line}
+        {"role": "user", "content": first_line + query}
     ]
+    print(messages)
     model_inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to("cuda")
-    generated_ids = llm.generate(model_inputs, max_new_tokens=512, do_sample=True)
+    print('Start Inference')
+    generated_ids = llm.generate(model_inputs, max_new_tokens=50, do_sample=True)
+    print('Start detokenize')
     response = tokenizer.batch_decode(generated_ids)[0]
 
     # # Check if a GPU is available
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    response = response + f". Using device: {device}"
+    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # response = response + f". Using device: {device}"
     
     return response
 

requirements.txt

@@ -1 +1,2 @@
 transformers==4.36.0
+sentencepiece