Add generation code

app.py

@@ -2,13 +2,56 @@ import gradio as gr
 import os
 from huggingface_hub import login
 from dotenv import load_dotenv
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
+import torch
+from threading import Thread
 
 load_dotenv()
 hf_token = os.getenv('HF_TOKEN', 'your-key-if-not-using-env')
 login(hf_token, add_to_git_credential=True)
 
-def greet(name):
-    return "Hello " + name + "!!"
+MODEL = "m-a-p/OpenCodeInterpreter-DS-33B"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
+system_message = "You are a computer programmer that can translate python code to C++ in order to improve performance"
+
+def user_prompt_for(python):
+    return f"Rewrite this python code to C++. You must search for the maximum performance. \
+    Format your response in Markdown. This is the Code: \
+    \n\n\
+    {python}"
+
+def messages_for(python):
+    return [
+        {"role": "system", "content": system_message},
+        {"role": "user", "content": user_prompt_for(python)}
+    ]
+
+quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_quant_type="nf4"
+)
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL)
+tokenizer.pad_token = tokenizer.eos_token
+streamer = TextIteratorStreamer(tokenizer)
+
+model = AutoModelForCausalLM.from_pretrained(MODEL, device_map="auto", quantization_config=quant_config)
+
+cplusplus = None
+def translate(python):
+    inputs = tokenizer.apply_chat_template(messages_for(python), return_tensors="pt").to("cuda")
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=80)
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    cplusplus = ""
+    for chunk in streamer:
+        cplusplus += chunk
+        yield cplusplus
+
+    del inputs
+    torch.cuda.empty_cache()
+
+demo = gr.Interface(fn=translate, inputs="code", outputs="markdown")
 demo.launch()