feat: use spaces gpu

app.py

@@ -1,5 +1,6 @@
 import os
 import torch
+import spaces
 import gradio as gr
 from huggingface_hub import login
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -15,6 +16,7 @@ def get_llm(model_id):
     return model
 
 # 问答逻辑
+@spaces.GPU(duration=120)
 def retriever_qa(file, query):
     # 加载模型和分词器
     model_id = 'mistralai/Mistral-7B-Instruct-v0.2'
@@ -41,13 +43,17 @@ def retriever_qa(file, query):
 
         # Tokenize 输入
         model_inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(device)
+        print(f"Model Inputs: {model_inputs}")
         print('Start Inference')
         
         # 推理
-        generated_ids = llm.generate(model_inputs['input_ids'], max_new_tokens=50, do_sample=True)
-        
+        generated_ids = llm.generate(model_inputs, max_new_tokens=50, do_sample=True)
+        # generated_ids = llm.generate(input_ids=model_inputs['input_ids'], attention_mask=model_inputs['attention_mask'], max_new_tokens=50, do_sample=True)
+        print(f'Generated ids: {generated_ids}')
         # 解码输出
+        print('Start detokenize')
         response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        print(response)
         return response
 
     # 调用推理逻辑