Upload 3 files

app.py

@@ -1,42 +1,32 @@
 import os
 import gradio as gr
 from huggingface_hub import login
-from transformers import pipeline
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Load the gated model
 #model_name = "RickyDeSkywalker/TheoremLlama"
-model_name = "unsloth/Llama-3.2-1B-Instruct"
-device = "cuda" if torch.cuda.is_available() else "cpu"
 #model_name = "internlm/internlm2-math-plus-7b"
 #model_name = "deepseek-ai/DeepSeek-Prover-V1.5-RL"
+model_name = "unsloth/Llama-3.2-1B-Instruct"
+
 HF_TOKEN = os.environ.get("HF_TOKEN")
 #login(HF_TOKEN)
 
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-# Set `torch_dtype=torch.float16` to load model in float16, otherwise it will be loaded as float32 and might cause OOM Error.
 model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16, device_map="auto").eval()
-terminators = [tokenizer.eos_token_id, 
-               tokenizer.convert_tokens_to_ids("<|eot_id|>"), 
-               tokenizer.convert_tokens_to_ids("<|reserved_special_token_26|>")]
-
-
-#generator = pipeline('text-generation', model=model_name, trust_remote_code=True, token=HF_TOKEN)
 
 # Function for generating Lean 4 code
 @torch.inference_mode()
-def generate_lean4_code(prompt):
-    #result = generator(prompt, max_length=1000, num_return_sequences=1)
-    #return result[0]['generated_text']
-    #response, history = model.chat(tokenizer, prompt, history=[], meta_instruction="")
-    #print(response, history)
-    #return response
+def generate_lean4_code(prompt: str):
+    terminators = [tokenizer.eos_token_id, 
+               tokenizer.convert_tokens_to_ids("<|eot_id|>"), 
+               tokenizer.convert_tokens_to_ids("<|reserved_special_token_26|>")]
     chat = [
         {"role": "system", "content": "You are a Lean4 expert who can write good Lean4 code based on natural language mathematical theorem and proof"},
         {"role": "user", "content": prompt},
     ]
-    input_ids = tokenizer.apply_chat_template(chat, return_tensors="pt").to(device)
+    input_ids = tokenizer.apply_chat_template(chat, return_tensors="pt").to(model.device)
     results = model.generate(input_ids, max_new_tokens=1024, eos_token_id=terminators, do_sample=True, temperature=0.85, top_p=0.9)
     result_str = tokenizer.decode(results[0], skip_special_tokens=True)
     return result_str
@@ -52,8 +42,7 @@ interface = gr.Interface(
         placeholder="Prove that the sum of two even numbers is even.",
         lines=4
     ),
-    #outputs=gr.Code(label="Generated Lean 4 Code", language="lean"),
-    outputs=gr.Code(label="Generated Lean 4 Code"),
+    outputs=gr.Code(label="Generated Lean 4 Code"), # language="lean" unsupported in Gradio...
     title=title,
     description=description
 )