Upload app.py

app.py

@@ -0,0 +1,51 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import transformers
+import torch
+import gradio as gr
+
+#Setting device to cuda
+torch.set_default_device("cuda")
+
+
+model = "deepapaikar/llama_mistral"
+
+pipeline = transformers.pipeline(
+    "text-generation",
+    model=model,
+    torch_dtype=torch.float16,
+    device_map="auto",
+)
+
+tokenizer = AutoTokenizer.from_pretrained(model)
+
+
+def predict_answer(question, token=25):
+    
+    messages = [{"role": "user", "content": f"{question}"}]
+    
+    
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    
+    
+    outputs = pipeline(prompt, max_new_tokens=token, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
+    
+    return outputs[0]["generated_text"]
+
+
+
+def gradio_predict(question, token):
+    answer = predict_answer(question, token)
+    return answer
+
+# Define the Gradio interface
+iface = gr.Interface(
+    fn=gradio_predict,
+    inputs=[gr.Textbox(label="Question", placeholder="e.g. What are the colors of the bus in the image?", scale=4),
+           gr.Slider(2, 100, value=25, label="Token Count", info="Choose between 2 and 100")],
+    outputs=gr.TextArea(label="Answer"),
+    title="KatzBot",
+    description="Llama - Mistral Merge",
+)
+
+# Launch the app
+iface.queue().launch(debug=True)