feat: add mistral

app.py

@@ -1,8 +1,12 @@
-import torch
+import spaces
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Check if a GPU is available
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
+model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2", device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
+
+# # Check if a GPU is available
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# print(f"Using device: {device}")
 
 import gradio as gr
 
@@ -13,8 +17,15 @@ import warnings
 warnings.warn = warn
 warnings.filterwarnings('ignore')
 
+def get_llm():
+    model_id = 'mistralai/Mistral-7B-Instruct-v0.2'
+    model = AutoModelForCausalLM.from_pretrained(model_id, device_map='auto')
+    model.to('cuda')
+    return model
+
+@spaces.GPU
 def retriever_qa(file, query):
-    # llm = get_llm()
+    llm = get_llm()
     # retriever_obj = retriever(file)
     # qa = RetrievalQA.from_chain_type(llm=llm, 
     #                                 chain_type="stuff", 
@@ -23,8 +34,22 @@ def retriever_qa(file, query):
     # response = qa.invoke(query)
     with open(file, 'r') as f:
         first_line = f.readline()
+
+    messages = [
+        {"role": "user", "content": first_line}
+    ]
+
+    model_inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to("cuda")
+
+    generated_ids = llm.generate(model_inputs, max_new_tokens=100, do_sample=True)
+    # tokenizer.batch_decode(generated_ids)[0]
     
-    response = first_line + query
+    response = tokenizer.batch_decode(generated_ids)[0]
+
+    # # Check if a GPU is available
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # print(f"Using device: {device}")
+    response = response + f". Using device: {device}"
     
     return response