Update app.py

app.py

@@ -1,14 +1,49 @@
-import gradio as gr
+import os
 import spaces
 import torch
+import gradio as gr
+from huggingface_hub import snapshot_download, login
+from transformers.utils import move_cache
+
+
+
+
+LLM_MODEL_DIR = '/model'
+LLM_MODEL_ID = "mistral-community/Mistral-7B-v0.2"
+LLM_MODEL_REVISION = 'main'
+os.makedirs(LLM_MODEL_DIR, exist_ok=True)
+snapshot_download(LLM_MODEL_ID, revision=LLM_MODEL_REVISION, local_dir=LLM_MODEL_DIR) #, token=HF_TOKEN)
+move_cache()
+
+
+# cpu
 
 zero = torch.Tensor([0]).cuda()
 print(zero.device) # <-- 'cpu' 🤔
 
+
+# gpu
+
 @spaces.GPU
-def greet(n):
-    print(zero.device) # <-- 'cuda:0' 🤗
-    return f"Hello {zero + n} Tensor"
+def greet(user):
+    # print(zero.device) # <-- 'cuda:0' 🤗
+    from vllm import SamplingParams, LLM
+
+    model = LLM(LLM_MODEL_DIR)
+    sampling_params = dict(
+            temperature = 0.3,
+            ignore_eos = False,
+            max_tokens = int(512 * 2)
+    )
+    sampling_params = SamplingParams(**sampling_params)
+    
+    prompts = [user]
+    model_outputs = model.generate(prompts, sampling_params)
+    generations = []
+    for output in model_outputs:
+        for outputs in output.outputs:
+            generations.append(outputs.text)
+    return generations[0]
 
 demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
 demo.launch()