test1

app.py

@@ -1,8 +1,39 @@
 
 import gradio as gr
+#import peft
+import transformers
 
-def greet(name):
-    return "Hello " + name + "!!"
+device = "cpu"
+is_peft = False
+model_id = "treadon/promt-fungineer-355M"
+# if is_peft:
+#     config = peft.PeftConfig.from_pretrained(model_id)
+#     model = transformers.AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, low_cpu_mem_usage=True)
+#     tokenizer = transformers.AutoTokenizer.from_pretrained(config.base_model_name_or_path)
+#     model = peft.PeftModel.from_pretrained(model, model_id)
+# else:
+model = transformers.AutoModelForCausalLM.from_pretrained(model_id, low_cpu_mem_usage=True)
+tokenizer = transformers.AutoTokenizer.from_pretrained("gpt2")
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+def generate_text(prompt):
+    
+    if not prompt.startswith("BRF:"):
+        prompt = "BRF: " + prompt
+
+    model.eval()
+    # SOFT SAMPLE
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    samples = []
+    try:
+        for i in range(1):
+            outputs = model.generate(**inputs, max_length=256, do_sample=True, top_k=100, top_p=0.95, temperature=0.85, num_return_sequences=4, pad_token_id=tokenizer.eos_token_id)
+            for output in outputs:
+                sample = tokenizer.decode(output, skip_special_tokens=True)
+                samples.append(sample)
+    except Exception as e:
+        print(e)
+
+    return samples
+
+iface = gr.Interface(fn=generate_text, inputs="text", outputs=("text","text","text","text") )
 iface.launch()