initial commit

app.py

@@ -0,0 +1,52 @@
+import gradio as gr
+import spaces
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+from huggingface_hub import login
+import os
+
+access_token = os.getenv('HF_TOKEN')
+login(access_token)
+
+model_id = "google/gemma-2-9b-it"
+tokenizer = None
+model = None
+
+@spaces.GPU
+def load_model():
+    global tokenizer, model
+    print("Model loading started")
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+    )
+    print("Model loading completed. Device of the model:", model.device)
+
+load_model()
+
+@spaces.GPU
+def ask(prompt):
+    if not prompt:
+        return {"error": "Prompt is missing"}
+
+    print("Device of the model:", model.device)
+    messages = [
+        {"role": "user", "content": f"{prompt}"},
+    ]
+    print("Messages:", messages)
+    print("Tokenizer process started")
+    input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt", return_dict=True).to("cuda")
+    print("Tokenizer process completed")
+
+    print("Model process started")
+    outputs = model.generate(**input_ids, max_new_tokens=256)
+
+    print("Tokenizer decode process started")
+    answer = tokenizer.decode(outputs[0]).split("<end_of_turn>")[1].strip()
+
+    return answer
+
+demo = gr.Interface(fn=ask, inputs=gr.Textbox(lines=2, placeholder="Enter your prompt here..."), outputs=gr.Textbox())
+demo.launch()

requirements.txt

@@ -0,0 +1,8 @@
+gradio
+transformers
+torch
+huggingface_hub
+requests
+gradio_client
+fastapi
+uvicorn