Upload app.py

app.py

@@ -12,30 +12,28 @@ DEFAULT_MAX_NEW_TOKENS = 1024
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 ACCESS_TOKEN = os.getenv("HF_TOKEN", "")
 
+model_id = "meta-llama/Llama-2-13b-chat"
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    trust_remote_code=True,
+    token=ACCESS_TOKEN)
+tokenizer = AutoTokenizer.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    token=ACCESS_TOKEN)
+tokenizer.use_default_system_prompt = False
+
+
 @spaces.GPU
 def generate(
-    model: str,
     message: str,
     system_prompt: str,
     max_new_tokens: int = 1024,
     temperature: float = 0.01,
     top_p: float = 0.01,
 ) -> Iterator[str]:
-
-    model_id = model
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16,
-        device_map="auto",
-        trust_remote_code=True,
-        token=ACCESS_TOKEN)
-    tokenizer = AutoTokenizer.from_pretrained(
-        model_id,
-        trust_remote_code=True,
-        token=ACCESS_TOKEN)
-    tokenizer.use_default_system_prompt = False
-
-
     conversation = []
     if system_prompt:
         conversation.append({"role": "system", "content": system_prompt})
@@ -75,7 +73,6 @@ def generate(
 chat_interface = gr.Interface(
     fn=generate,
     inputs=[
-        gr.Textbox(lines=1, placeholder="Model", label="Model name"),
         gr.Textbox(lines=2, placeholder="Prompt", label="Prompt"),
     ],
     outputs="text",