Update app.py

app.py

@@ -1,20 +1,12 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from vllm import LLM, SamplingParams
 import gradio as gr
+from llama_cpp import Llama
 
-model_name = "aws-prototyping/MegaBeam-Mistral-7B-512k"
+MODEL_PATH = "model.gguf"  # downloaded in advance
 
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    trust_remote_code=True
-)
+llm = Llama(model_path=MODEL_PATH, n_ctx=8192, n_threads=4)
 
-def chat(prompt: str):
-    llm = LLM(model=model_name)
-    sampling = SamplingParams(temperature=0.7, max_tokens=512)
-    outputs = llm.generate([prompt], sampling)
-    return outputs[0].outputs[0].text
+def chat(prompt):
+    response = llm(prompt, max_tokens=512, temperature=0.7)
+    return response["choices"][0]["text"]
 
-iface = gr.Interface(fn=chat, inputs="text", outputs="text")
-iface.launch()
+gr.Interface(fn=chat, inputs="text", outputs="text", title="MegaBeam Mistral 512K - GGUF").launch()